From 52997da2051966e63c2ae480be3f33e7a4e7921d Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 26 May 2026 08:32:53 +0200 Subject: [PATCH 01/12] Add HNSW layered hierarchy --- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu | 11 +- .../ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h | 30 +- cpp/include/cuvs/neighbors/hnsw.hpp | 18 +- cpp/src/neighbors/detail/hnsw.hpp | 1125 ++++++++++++++++- cpp/tests/neighbors/ann_hnsw_ace.cuh | 257 ++++ .../ann_hnsw_ace/test_float_uint32_t.cu | 10 + .../ann_hnsw_ace/test_half_uint32_t.cu | 20 + .../ann_hnsw_ace/test_int8_t_uint32_t.cu | 20 + .../ann_hnsw_ace/test_uint8_t_uint32_t.cu | 20 + examples/cpp/CMakeLists.txt | 2 + examples/cpp/src/hnsw_ace_layered_example.cu | 278 ++++ 11 files changed, 1723 insertions(+), 68 deletions(-) create mode 100644 examples/cpp/src/hnsw_ace_layered_example.cu diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu index 26028b6d98..b535ced763 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu @@ -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" @@ -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().c_str()); } @@ -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 diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index 2f0c54e1bd..0104ccdb3e 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/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 @@ -9,10 +9,26 @@ #include #include +#include #include 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 class cuvs_cagra_hnswlib : public algo, public algo_gpu { public: @@ -130,6 +146,18 @@ void cuvs_cagra_hnswlib::set_search_param(const search_param_base& para template void cuvs_cagra_hnswlib::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."); diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index fb726fed71..5800e05a96 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -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 { @@ -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. diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 4914a0fa1b..44dc4e688e 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -23,15 +23,28 @@ #include #include +#include +#include +#include +#include +#include +#include #include #include #include +#include #include #include +#include +#include #include +#include +#include #include #include #include +#include +#include namespace cuvs::neighbors::hnsw::detail { @@ -166,6 +179,9 @@ struct index_impl : index { RAFT_LOG_INFO("Loading HNSW index from disk: %s", filepath.c_str()); try { + RAFT_EXPECTS(this->hierarchy() != HnswHierarchy::GPU_LAYERED_ON_DISK, + "Layered HNSW indexes must be loaded with hnsw::deserialize so a local dataset " + "can be provided through index_params.dataset_path."); appr_alg_ = std::make_unique::type>>( space_.get(), filepath); if (this->hierarchy() == HnswHierarchy::NONE) { appr_alg_->base_layer_only = true; } @@ -341,6 +357,796 @@ void all_neighbors_graph(raft::resources const& res, } } +struct hnsw_level_plan { + size_t n_rows = 0; + std::vector hist; + // Bucket-end offsets after order construction. Level L starts at offsets[L - 1]. + std::vector offsets; + std::vector order; + std::vector order_bw; + std::vector levels; + + [[nodiscard]] auto max_level() const -> int + { + return hist.empty() ? 0 : static_cast(hist.size() - 1); + } + + [[nodiscard]] auto promoted_count() const -> size_t + { + return hist.empty() ? 0 : n_rows - hist[0]; + } +}; + +inline auto hnsw_upper_layer_degree(size_t row_count, size_t M) -> size_t +{ + return row_count > M ? M : row_count > 0 ? row_count - 1 : 0; +} + +inline auto elapsed_ms_since(std::chrono::steady_clock::time_point start) -> int64_t +{ + return std::chrono::duration_cast(std::chrono::steady_clock::now() - + start) + .count(); +} + +inline auto to_gib(size_t bytes) -> double +{ + return static_cast(bytes) / (1024.0 * 1024.0 * 1024.0); +} + +inline auto throughput_gib_per_s(size_t bytes, int64_t elapsed_ms) -> double +{ + return elapsed_ms > 0 ? to_gib(bytes) / (elapsed_ms / 1000.0) : 0.0; +} + +inline auto make_hnsw_level_plan_from_levels(size_t n_rows, + std::vector&& levels, + bool build_reverse_order, + const char* log_prefix = nullptr) -> hnsw_level_plan +{ + hnsw_level_plan plan; + plan.n_rows = n_rows; + plan.levels = std::move(levels); + RAFT_EXPECTS(n_rows > 0, "HNSW hierarchy requires at least one row"); + RAFT_EXPECTS(plan.levels.size() == n_rows, + "HNSW level count (%zu) must match row count (%zu)", + plan.levels.size(), + n_rows); + + for (auto level : plan.levels) { + while (static_cast(level) >= plan.hist.size()) { + plan.hist.push_back(0); + } + plan.hist[level]++; + } + + plan.offsets.resize(plan.hist.size() + 1, 0); + for (size_t i = 0; i < plan.hist.size() - 1; ++i) { + plan.offsets[i + 1] = plan.offsets[i] + plan.hist[i]; + if (log_prefix != nullptr) { + RAFT_LOG_INFO("%s%zu : %zu", log_prefix, i + 1, n_rows - plan.offsets[i + 1]); + } + } + + plan.order.resize(n_rows); + if (build_reverse_order) { plan.order_bw.resize(n_rows); } + for (size_t i = 0; i < n_rows; ++i) { + const auto level = static_cast(plan.levels[i]); + if (build_reverse_order) { plan.order_bw[i] = plan.offsets[level]; } + plan.order[plan.offsets[level]++] = i; + } + + return plan; +} + +template +auto make_random_hnsw_level_plan(size_t n_rows, HnswAlgo& appr_algo, const char* log_prefix) + -> hnsw_level_plan +{ + std::vector levels(n_rows); + for (int64_t i = 0; i < static_cast(n_rows); ++i) { + const auto pt_level = appr_algo.getRandomLevel(appr_algo.mult_); + RAFT_EXPECTS(pt_level <= std::numeric_limits::max(), + "HNSW serialization only supports levels up to %u", + static_cast(std::numeric_limits::max())); + levels[i] = static_cast(pt_level); + } + + return make_hnsw_level_plan_from_levels(n_rows, std::move(levels), true, log_prefix); +} + +template +void build_hnsw_upper_layer_graphs( + raft::resources const& res, + raft::host_matrix_view promoted_dataset, + const hnsw_level_plan& plan, + size_t M, + cuvs::distance::DistanceType metric, + Callback&& callback) +{ + static_assert(std::is_same_v, + "HNSW upper-layer graph construction expects uint32_t neighbor ids"); + + const auto dim = static_cast(promoted_dataset.extent(1)); + for (size_t pt_level = 1; pt_level < plan.hist.size(); ++pt_level) { + const auto start_idx = plan.offsets[pt_level - 1]; + const auto removed_rows = start_idx - plan.offsets[0]; + const auto row_count = plan.n_rows - start_idx; + const auto neighbor_size = hnsw_upper_layer_degree(row_count, M); + + RAFT_LOG_INFO("Compute hierarchy neighbors level %zu", pt_level); + auto host_neighbors = raft::make_host_matrix( + static_cast(row_count), static_cast(neighbor_size)); + + if (neighbor_size > 0) { + auto layer_dataset_view = raft::make_host_matrix_view( + promoted_dataset.data_handle() + removed_rows * dim, + static_cast(row_count), + promoted_dataset.extent(1)); + all_neighbors_graph(res, layer_dataset_view, host_neighbors.view(), metric); + } + + callback(pt_level, start_idx, host_neighbors); + } +} + +struct layered_hnsw_layer_info { + size_t level = 0; + size_t row_count = 0; + size_t degree = 0; + size_t node_offset = 0; + size_t link_offset = 0; +}; + +struct layered_hnsw_file_metadata { + size_t n_rows = 0; + size_t dim = 0; + size_t M = 0; + size_t maxM = 0; + size_t maxM0 = 0; + size_t ef_construction = 0; + size_t base_degree = 0; + size_t levels_bytes = 0; + size_t base_link_row_bytes = 0; + size_t base_links_bytes = 0; + size_t upper_nodes_count = 0; + size_t upper_nodes_bytes = 0; + size_t upper_link_row_bytes = 0; + size_t upper_links_bytes = 0; + double mult = 0.0; + int maxlevel = 0; + int enterpoint_node = 0; + std::vector layers; +}; + +struct layered_hnsw_file_header { + char magic[32]; + uint32_t version; + uint32_t reserved; + uint64_t metadata_size; + uint64_t metadata_offset; +}; + +constexpr const char* layered_hnsw_magic = "CUVS_HNSW_LAYERED"; +constexpr uint32_t layered_hnsw_version = 1; +constexpr size_t layered_hnsw_alignment = 64; + +inline auto align_up(size_t value, size_t alignment) -> size_t +{ + return ((value + alignment - 1) / alignment) * alignment; +} + +template +auto layered_dtype_name() -> const char* +{ + if constexpr (std::is_same_v) { + return "float32"; + } else if constexpr (std::is_same_v) { + return "float16"; + } else if constexpr (std::is_same_v) { + return "uint8"; + } else if constexpr (std::is_same_v) { + return "int8"; + } else { + return "unknown"; + } +} + +inline auto metric_name(cuvs::distance::DistanceType metric) -> const char* +{ + switch (metric) { + case cuvs::distance::DistanceType::L2Expanded: return "L2Expanded"; + case cuvs::distance::DistanceType::InnerProduct: return "InnerProduct"; + default: return "Unknown"; + } +} + +template +auto make_layered_hnsw_metadata_json(const layered_hnsw_file_metadata& metadata, + cuvs::distance::DistanceType metric) -> std::string +{ + std::stringstream os; + os << "{\n"; + os << " \"format\": \"cuvs_hnsw_layered\",\n"; + os << " \"version\": 1,\n"; + os << " \"n_rows\": " << metadata.n_rows << ",\n"; + os << " \"dim\": " << metadata.dim << ",\n"; + os << " \"dtype\": \"" << layered_dtype_name() << "\",\n"; + os << " \"metric\": \"" << metric_name(metric) << "\",\n"; + os << " \"row_order\": \"original_id\",\n"; + os << " \"M\": " << metadata.M << ",\n"; + os << " \"maxM\": " << metadata.maxM << ",\n"; + os << " \"maxM0\": " << metadata.maxM0 << ",\n"; + os << " \"ef_construction\": " << metadata.ef_construction << ",\n"; + os << " \"mult\": " << metadata.mult << ",\n"; + os << " \"maxlevel\": " << metadata.maxlevel << ",\n"; + os << " \"enterpoint_node\": " << metadata.enterpoint_node << ",\n"; + os << " \"base_degree\": " << metadata.base_degree << ",\n"; + os << " \"levels_bytes\": " << metadata.levels_bytes << ",\n"; + os << " \"base_link_row_bytes\": " << metadata.base_link_row_bytes << ",\n"; + os << " \"base_links_bytes\": " << metadata.base_links_bytes << ",\n"; + os << " \"upper_nodes_count\": " << metadata.upper_nodes_count << ",\n"; + os << " \"upper_nodes_bytes\": " << metadata.upper_nodes_bytes << ",\n"; + os << " \"upper_link_row_bytes\": " << metadata.upper_link_row_bytes << ",\n"; + os << " \"upper_links_bytes\": " << metadata.upper_links_bytes << ",\n"; + os << " \"layers\": [\n"; + for (size_t i = 0; i < metadata.layers.size(); ++i) { + const auto& layer = metadata.layers[i]; + os << " {\"level\": " << layer.level << ", \"row_count\": " << layer.row_count + << ", \"degree\": " << layer.degree << ", \"node_offset\": " << layer.node_offset + << ", \"link_offset\": " << layer.link_offset << "}"; + os << (i + 1 == metadata.layers.size() ? "\n" : ",\n"); + } + os << " ]\n"; + os << "}\n"; + return os.str(); +} + +inline auto json_find_key(const std::string& json, const std::string& key, size_t start = 0) + -> size_t +{ + const auto pos = json.find("\"" + key + "\"", start); + RAFT_EXPECTS( + pos != std::string::npos, "Cannot find key '%s' in layered HNSW metadata", key.c_str()); + return pos; +} + +inline auto json_parse_size(const std::string& json, const std::string& key) -> size_t +{ + auto pos = json_find_key(json, key); + auto colon = json.find(':', pos); + RAFT_EXPECTS(colon != std::string::npos, "Malformed JSON near key '%s'", key.c_str()); + auto begin = json.find_first_of("0123456789", colon + 1); + auto end = json.find_first_not_of("0123456789", begin); + RAFT_EXPECTS( + begin != std::string::npos, "Malformed integer JSON value for key '%s'", key.c_str()); + return static_cast(std::stoull(json.substr(begin, end - begin))); +} + +inline auto json_parse_double(const std::string& json, const std::string& key) -> double +{ + auto pos = json_find_key(json, key); + auto colon = json.find(':', pos); + RAFT_EXPECTS(colon != std::string::npos, "Malformed JSON near key '%s'", key.c_str()); + auto begin = json.find_first_of("0123456789-.", colon + 1); + auto end = json.find_first_not_of("0123456789-.eE+", begin); + RAFT_EXPECTS(begin != std::string::npos, "Malformed double JSON value for key '%s'", key.c_str()); + return std::stod(json.substr(begin, end - begin)); +} + +inline auto json_parse_layer_field(const std::string& object, const std::string& key) -> size_t +{ + return json_parse_size(object, key); +} + +inline auto parse_layered_hnsw_metadata(const std::string& json) -> layered_hnsw_file_metadata +{ + layered_hnsw_file_metadata metadata; + metadata.n_rows = json_parse_size(json, "n_rows"); + metadata.dim = json_parse_size(json, "dim"); + metadata.M = json_parse_size(json, "M"); + metadata.maxM = json_parse_size(json, "maxM"); + metadata.maxM0 = json_parse_size(json, "maxM0"); + metadata.ef_construction = json_parse_size(json, "ef_construction"); + metadata.mult = json_parse_double(json, "mult"); + metadata.maxlevel = static_cast(json_parse_size(json, "maxlevel")); + metadata.enterpoint_node = static_cast(json_parse_size(json, "enterpoint_node")); + metadata.base_degree = json_parse_size(json, "base_degree"); + metadata.levels_bytes = json_parse_size(json, "levels_bytes"); + metadata.base_link_row_bytes = json_parse_size(json, "base_link_row_bytes"); + metadata.base_links_bytes = json_parse_size(json, "base_links_bytes"); + metadata.upper_nodes_count = json_parse_size(json, "upper_nodes_count"); + metadata.upper_nodes_bytes = json_parse_size(json, "upper_nodes_bytes"); + metadata.upper_link_row_bytes = json_parse_size(json, "upper_link_row_bytes"); + metadata.upper_links_bytes = json_parse_size(json, "upper_links_bytes"); + + auto layers_pos = json_find_key(json, "layers"); + auto array_open = json.find('[', layers_pos); + auto array_end = json.find(']', array_open); + RAFT_EXPECTS(array_open != std::string::npos && array_end != std::string::npos, + "Malformed layers array in layered HNSW metadata"); + + size_t pos = array_open; + while (true) { + auto object_open = json.find('{', pos); + if (object_open == std::string::npos || object_open > array_end) { break; } + auto object_end = json.find('}', object_open); + RAFT_EXPECTS(object_end != std::string::npos && object_end < array_end, + "Malformed layer object in layered HNSW metadata"); + const auto object = json.substr(object_open, object_end - object_open + 1); + metadata.layers.push_back({json_parse_layer_field(object, "level"), + json_parse_layer_field(object, "row_count"), + json_parse_layer_field(object, "degree"), + json_parse_layer_field(object, "node_offset"), + json_parse_layer_field(object, "link_offset")}); + pos = object_end + 1; + } + return metadata; +} + +struct npy_file { + cuvs::util::file_descriptor fd; + size_t header_size = 0; + std::vector shape; +}; + +inline auto open_npy_file(const std::string& path) -> npy_file +{ + std::ifstream stream(path, std::ios::binary); + RAFT_EXPECTS(stream.good(), "Failed to open numpy file: %s", path.c_str()); + auto header = raft::detail::numpy_serializer::read_header(stream); + auto header_size = static_cast(stream.tellg()); + auto fd = cuvs::util::file_descriptor(path, O_RDONLY); + return {std::move(fd), header_size, header.shape}; +} + +inline auto ends_with(const std::string& value, const std::string& suffix) -> bool +{ + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +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); +} + +template +inline auto open_layered_dataset_file(const std::string& path) -> npy_file +{ + if (std::filesystem::path(path).extension() == ".npy") { return open_npy_file(path); } + + std::ifstream stream(path, std::ios::binary); + RAFT_EXPECTS(stream.good(), "Failed to open dataset file: %s", path.c_str()); + std::array header{}; + stream.read(reinterpret_cast(header.data()), sizeof(uint32_t) * header.size()); + RAFT_EXPECTS(stream.gcount() == static_cast(sizeof(uint32_t) * header.size()), + "Failed to read ANN benchmark binary dataset header: %s", + path.c_str()); + + const auto ext = std::filesystem::path(path).extension().string(); + if constexpr (std::is_same_v) { + RAFT_EXPECTS(ext == ".fbin" && !ends_with(path, ".fp16.fbin"), + "Expected a .fbin dataset for float layered HNSW load: %s", + path.c_str()); + } else if constexpr (std::is_same_v) { + RAFT_EXPECTS(ext == ".f16bin" || ends_with(path, ".fp16.fbin"), + "Expected a .f16bin or .fp16.fbin dataset for half layered HNSW load: %s", + path.c_str()); + } else if constexpr (std::is_same_v) { + RAFT_EXPECTS( + ext == ".u8bin", "Expected a .u8bin dataset for uint8 layered HNSW load: %s", path.c_str()); + } else if constexpr (std::is_same_v) { + RAFT_EXPECTS( + ext == ".i8bin", "Expected a .i8bin dataset for int8 layered HNSW load: %s", path.c_str()); + } + + auto fd = cuvs::util::file_descriptor(path, O_RDONLY); + return {std::move(fd), sizeof(uint32_t) * header.size(), {header[0], header[1]}}; +} + +template +void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index& index_, + const cuvs::util::file_descriptor& output_fd, + size_t output_offset, + size_t base_link_row_bytes, + size_t maxM0) +{ + const auto total_start_time = std::chrono::steady_clock::now(); + const auto& graph_fd_opt = index_.graph_fd(); + const auto& mapping_fd_opt = index_.mapping_fd(); + RAFT_EXPECTS(graph_fd_opt.has_value() && graph_fd_opt->is_valid(), + "Graph file descriptor is not available"); + RAFT_EXPECTS(mapping_fd_opt.has_value() && mapping_fd_opt->is_valid(), + "Mapping file descriptor is not available"); + + const auto graph_path = graph_fd_opt->get_path(); + const auto mapping_path = mapping_fd_opt->get_path(); + RAFT_EXPECTS(!graph_path.empty(), "Unable to get path from graph file descriptor"); + RAFT_EXPECTS(!mapping_path.empty(), "Unable to get path from mapping file descriptor"); + + auto graph_npy = open_npy_file(graph_path); + auto mapping_npy = open_npy_file(mapping_path); + + RAFT_EXPECTS(graph_npy.shape.size() == 2, "Graph file should be 2D"); + RAFT_EXPECTS(mapping_npy.shape.size() == 1, "Mapping file should be 1D"); + const auto n_rows = graph_npy.shape[0]; + const auto degree = graph_npy.shape[1]; + RAFT_EXPECTS(mapping_npy.shape[0] == n_rows, + "Mapping elements (%zu) != graph rows (%zu)", + mapping_npy.shape[0], + n_rows); + RAFT_EXPECTS(n_rows == static_cast(index_.size()), + "Graph rows (%zu) != index size (%zu)", + n_rows, + static_cast(index_.size())); + RAFT_EXPECTS(degree == static_cast(index_.graph_degree()), + "Graph degree (%zu) != index graph degree (%zu)", + degree, + static_cast(index_.graph_degree())); + RAFT_EXPECTS(degree > 0, "Graph degree must be nonzero"); + RAFT_EXPECTS( + degree <= maxM0, "Base graph degree (%zu) must not exceed HNSW maxM0 (%zu)", degree, maxM0); + RAFT_EXPECTS(base_link_row_bytes >= sizeof(hnswlib::linklistsizeint) + maxM0 * sizeof(IdxT), + "Base link row size is too small"); + + RAFT_LOG_INFO("HNSW remap: loading ACE reordered-to-original mapping (%zu rows)", n_rows); + const auto mapping_start_time = std::chrono::steady_clock::now(); + const auto mapping_bytes = n_rows * sizeof(IdxT); + std::vector reordered_to_original(n_rows); + std::vector original_to_reordered(n_rows); + cuvs::util::read_large_file( + mapping_npy.fd, reordered_to_original.data(), mapping_bytes, mapping_npy.header_size); + for (size_t reordered_id = 0; reordered_id < n_rows; ++reordered_id) { + const auto original_id = static_cast(reordered_to_original[reordered_id]); + RAFT_EXPECTS(original_id < n_rows, + "Invalid original id %zu in ACE dataset mapping at row %zu", + original_id, + reordered_id); + original_to_reordered[original_id] = static_cast(reordered_id); + } + const auto mapping_elapsed_ms = elapsed_ms_since(mapping_start_time); + RAFT_LOG_INFO("HNSW remap: mapping loaded in %ld ms (%.2f GiB, %.2f GiB/s)", + mapping_elapsed_ms, + to_gib(mapping_bytes), + throughput_gib_per_s(mapping_bytes, mapping_elapsed_ms)); + + const size_t graph_row_bytes = degree * sizeof(IdxT); + const size_t target_batch_bytes = 64 * 1024 * 1024; + const size_t batch_size = std::max(1, target_batch_bytes / base_link_row_bytes); + const size_t max_gap_rows = std::max(1, (128 * 1024) / graph_row_bytes); + const size_t max_span_rows = std::max(1, (8 * 1024 * 1024) / graph_row_bytes); + + RAFT_LOG_INFO( + "HNSW remap: n_rows=%zu degree=%zu maxM0=%zu batch_size=%zu max_gap=%zu rows max_span=%zu rows", + n_rows, + degree, + maxM0, + batch_size, + max_gap_rows, + max_span_rows); + + auto span_buffer = raft::make_host_matrix(static_cast(max_span_rows), + static_cast(degree)); + std::vector output_buffer(batch_size * base_link_row_bytes); + std::vector source_rows(batch_size); + std::vector order(batch_size); + + const auto graph_start_time = std::chrono::steady_clock::now(); + size_t spans_read = 0; + size_t bytes_read = 0; + size_t bytes_written = 0; + + for (size_t batch_start = 0; batch_start < n_rows; batch_start += batch_size) { + const auto current_batch_size = std::min(batch_size, n_rows - batch_start); + std::fill( + output_buffer.begin(), output_buffer.begin() + current_batch_size * base_link_row_bytes, 0); + for (size_t i = 0; i < current_batch_size; ++i) { + source_rows[i] = original_to_reordered[batch_start + i]; + order[i] = i; + } + std::sort(order.begin(), order.begin() + current_batch_size, [&](size_t lhs, size_t rhs) { + return source_rows[lhs] < source_rows[rhs]; + }); + + size_t sorted_pos = 0; + while (sorted_pos < current_batch_size) { + const auto span_start = static_cast(source_rows[order[sorted_pos]]); + auto span_end = span_start; + auto span_next = sorted_pos + 1; + while (span_next < current_batch_size) { + const auto next_row = static_cast(source_rows[order[span_next]]); + if (next_row - span_start + 1 > max_span_rows) { break; } + if (next_row > span_end + max_gap_rows + 1) { break; } + span_end = next_row; + ++span_next; + } + + const auto span_rows = span_end - span_start + 1; + const auto span_bytes = span_rows * graph_row_bytes; + cuvs::util::read_large_file(graph_npy.fd, + span_buffer.data_handle(), + span_bytes, + graph_npy.header_size + span_start * graph_row_bytes); + ++spans_read; + bytes_read += span_bytes; + + for (size_t i = sorted_pos; i < span_next; ++i) { + const auto dst_row = order[i]; + const auto source_row = static_cast(source_rows[dst_row]); + auto* src = span_buffer.data_handle() + (source_row - span_start) * degree; + auto* dst_row_ptr = output_buffer.data() + dst_row * base_link_row_bytes; + hnswlib::linklistsizeint list_count = static_cast(degree); + std::memcpy(dst_row_ptr, &list_count, sizeof(list_count)); + auto* dst = reinterpret_cast(dst_row_ptr + sizeof(hnswlib::linklistsizeint)); + for (size_t j = 0; j < degree; ++j) { + const auto neighbor = static_cast(src[j]); + RAFT_EXPECTS(neighbor < n_rows, + "Invalid reordered neighbor id %zu in ACE graph row %zu", + neighbor, + source_row); + dst[j] = reordered_to_original[neighbor]; + } + } + sorted_pos = span_next; + } + + const auto batch_bytes = current_batch_size * base_link_row_bytes; + cuvs::util::write_large_file(output_fd, + output_buffer.data(), + batch_bytes, + output_offset + batch_start * base_link_row_bytes); + bytes_written += batch_bytes; + } + + const auto graph_elapsed_ms = elapsed_ms_since(graph_start_time); + const auto total_elapsed_ms = elapsed_ms_since(total_start_time); + const auto read_gib = to_gib(bytes_read); + const auto written_gib = to_gib(bytes_written); + const auto ideal_read_bytes = n_rows * graph_row_bytes; + const auto amplification = + ideal_read_bytes == 0 ? 0.0 : static_cast(bytes_read) / ideal_read_bytes; + RAFT_LOG_INFO( + "HNSW remap: base links written in %ld ms: read %.2f GiB at %.2f GiB/s across %zu spans " + "(%.2fx amplification), wrote %.2f GiB at %.2f GiB/s", + graph_elapsed_ms, + read_gib, + throughput_gib_per_s(bytes_read, graph_elapsed_ms), + spans_read, + amplification, + written_gib, + throughput_gib_per_s(bytes_written, graph_elapsed_ms)); + RAFT_LOG_INFO("HNSW remap: completed in %ld ms total", total_elapsed_ms); +} + +template +auto serialize_to_layered_hnsw_from_disk( + raft::resources const& res, + const cuvs::neighbors::hnsw::index_params& params, + const cuvs::neighbors::cagra::index& index_, + raft::host_matrix_view dataset) -> std::string +{ + raft::common::nvtx::range fun_scope("hnsw::serialize_layered"); + const auto total_start_time = std::chrono::steady_clock::now(); + + RAFT_EXPECTS(index_.graph_fd().has_value() && index_.mapping_fd().has_value(), + "Layered HNSW serialization requires a disk-backed ACE graph and mapping."); + RAFT_EXPECTS(static_cast(dataset.extent(0)) == static_cast(index_.size()), + "Dataset rows (%zu) must match index size (%zu)", + static_cast(dataset.extent(0)), + static_cast(index_.size())); + RAFT_EXPECTS(static_cast(dataset.extent(1)) == static_cast(index_.dim()), + "Dataset cols (%zu) must match index dimensions (%zu)", + static_cast(dataset.extent(1)), + static_cast(index_.dim())); + + const auto graph_path = index_.graph_fd()->get_path(); + RAFT_EXPECTS(!graph_path.empty(), "Unable to get path from graph file descriptor"); + const auto ace_dir = std::filesystem::path(graph_path).parent_path(); + const auto artifact_file = ace_dir / "hnsw_index.cuvs"; + std::filesystem::create_directories(ace_dir); + + auto n_rows = static_cast(index_.size()); + auto dim = static_cast(index_.dim()); + auto graph_degree_int = static_cast(index_.graph_degree()); + + auto hnsw_index = + std::make_unique>(index_.dim(), index_.metric(), params.hierarchy); + auto appr_algo = std::make_unique::type>>( + hnsw_index->get_space(), 1, (graph_degree_int + 1) / 2, params.ef_construction); + + RAFT_LOG_INFO("Layered HNSW: generating hierarchy levels"); + const auto hierarchy_start_time = std::chrono::steady_clock::now(); + auto hierarchy = make_random_hnsw_level_plan(n_rows, *appr_algo, "Layered HNSW: Level "); + const auto hierarchy_elapsed_ms = elapsed_ms_since(hierarchy_start_time); + RAFT_LOG_INFO("Layered HNSW: hierarchy levels generated in %ld ms (max_level=%d, promoted=%zu)", + hierarchy_elapsed_ms, + hierarchy.max_level(), + hierarchy.promoted_count()); + + layered_hnsw_file_metadata metadata; + metadata.n_rows = n_rows; + metadata.dim = dim; + metadata.M = appr_algo->M_; + metadata.maxM = appr_algo->maxM_; + metadata.maxM0 = appr_algo->maxM0_; + metadata.ef_construction = appr_algo->ef_construction_; + metadata.mult = appr_algo->mult_; + metadata.maxlevel = hierarchy.max_level(); + metadata.enterpoint_node = static_cast(hierarchy.order.back()); + metadata.base_degree = static_cast(graph_degree_int); + metadata.levels_bytes = n_rows * sizeof(uint8_t); + metadata.base_link_row_bytes = appr_algo->size_links_level0_; + metadata.base_links_bytes = n_rows * metadata.base_link_row_bytes; + metadata.upper_link_row_bytes = appr_algo->size_links_per_element_; + size_t next_upper_node_offset = 0; + size_t next_upper_link_offset = 0; + for (size_t pt_level = 1; pt_level < hierarchy.hist.size(); pt_level++) { + auto start_idx = hierarchy.offsets[pt_level - 1]; + auto row_count = n_rows - start_idx; + auto layer_degree = hnsw_upper_layer_degree(row_count, appr_algo->M_); + metadata.layers.push_back( + {pt_level, row_count, layer_degree, next_upper_node_offset, next_upper_link_offset}); + next_upper_node_offset += row_count; + next_upper_link_offset += row_count; + } + metadata.upper_nodes_count = next_upper_node_offset; + metadata.upper_nodes_bytes = metadata.upper_nodes_count * sizeof(IdxT); + metadata.upper_links_bytes = next_upper_link_offset * metadata.upper_link_row_bytes; + + const auto metadata_json = make_layered_hnsw_metadata_json(metadata, index_.metric()); + layered_hnsw_file_header header{}; + std::strncpy(header.magic, layered_hnsw_magic, sizeof(header.magic) - 1); + header.version = layered_hnsw_version; + header.metadata_size = metadata_json.size(); + header.metadata_offset = sizeof(layered_hnsw_file_header); + + const auto payload_offset = + align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); + const auto levels_offset = payload_offset; + const auto base_links_offset = levels_offset + metadata.levels_bytes; + const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; + const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; + const auto final_file_size = upper_links_offset + metadata.upper_links_bytes; + + cuvs::util::file_descriptor artifact_fd(artifact_file.string(), O_CREAT | O_RDWR | O_TRUNC, 0644); + const auto fallocate_result = posix_fallocate(artifact_fd.get(), 0, final_file_size); + RAFT_EXPECTS(fallocate_result == 0, + "Failed to pre-allocate layered HNSW artifact %s: %s", + artifact_file.string().c_str(), + std::strerror(fallocate_result)); + cuvs::util::write_large_file(artifact_fd, &header, sizeof(header), 0); + cuvs::util::write_large_file( + artifact_fd, metadata_json.data(), metadata_json.size(), header.metadata_offset); + + const auto levels_start_time = std::chrono::steady_clock::now(); + cuvs::util::write_large_file( + artifact_fd, hierarchy.levels.data(), metadata.levels_bytes, levels_offset); + const auto levels_elapsed_ms = elapsed_ms_since(levels_start_time); + RAFT_LOG_INFO("Layered HNSW: wrote levels section in %ld ms (%.2f GiB, %.2f GiB/s)", + levels_elapsed_ms, + to_gib(metadata.levels_bytes), + throughput_gib_per_s(metadata.levels_bytes, levels_elapsed_ms)); + + RAFT_LOG_INFO("Layered HNSW: writing hnswlib-ready base links"); + const auto layer0_start_time = std::chrono::steady_clock::now(); + write_layered_base_links_from_disk( + index_, artifact_fd, base_links_offset, metadata.base_link_row_bytes, metadata.maxM0); + const auto layer0_elapsed_ms = elapsed_ms_since(layer0_start_time); + RAFT_LOG_INFO("Layered HNSW: base links written in %ld ms (%.2f GiB, %.2f GiB/s effective)", + layer0_elapsed_ms, + to_gib(metadata.base_links_bytes), + throughput_gib_per_s(metadata.base_links_bytes, layer0_elapsed_ms)); + + size_t upper_graph_bytes_written = 0; + if (hierarchy.hist.size() > 1) { + RAFT_LOG_INFO("Layered HNSW: gathering promoted vectors"); + const auto gather_start_time = std::chrono::steady_clock::now(); + auto host_query_set = + raft::make_host_matrix(static_cast(hierarchy.promoted_count()), dim); +#pragma omp parallel for + for (int64_t i = 0; i < static_cast(n_rows); i++) { + if (hierarchy.levels[i] > 0) { + const auto query_row = hierarchy.order_bw[i] - hierarchy.hist[0]; + auto* dst = host_query_set.data_handle() + query_row * dim; + std::copy(&dataset(i, 0), &dataset(i, 0) + dim, dst); + } + } + const auto gather_elapsed_ms = elapsed_ms_since(gather_start_time); + const auto gathered_bytes = hierarchy.promoted_count() * dim * sizeof(T); + RAFT_LOG_INFO("Layered HNSW: gathered promoted vectors in %ld ms (%.2f GiB copied, %.2f GiB/s)", + gather_elapsed_ms, + to_gib(gathered_bytes), + throughput_gib_per_s(gathered_bytes, gather_elapsed_ms)); + + auto promoted_dataset = raft::make_host_matrix_view( + host_query_set.data_handle(), host_query_set.extent(0), host_query_set.extent(1)); + const auto upper_layers_start_time = std::chrono::steady_clock::now(); + build_hnsw_upper_layer_graphs( + res, + promoted_dataset, + hierarchy, + appr_algo->M_, + index_.metric(), + [&](size_t pt_level, size_t start_idx, auto& host_neighbors) { + const auto& layer = metadata.layers[pt_level - 1]; + RAFT_LOG_INFO("Layered HNSW: writing upper layer %zu (%zu rows, degree %zu)", + layer.level, + layer.row_count, + layer.degree); + const auto layer_write_start_time = std::chrono::steady_clock::now(); + const size_t target_batch_bytes = 64 * 1024 * 1024; + const size_t row_bytes = sizeof(IdxT) + metadata.upper_link_row_bytes; + const size_t batch_size = std::max(1, target_batch_bytes / row_bytes); + std::vector node_buffer(batch_size); + std::vector link_buffer(batch_size * metadata.upper_link_row_bytes); + for (size_t batch_start = 0; batch_start < layer.row_count; batch_start += batch_size) { + const auto current_batch_size = std::min(batch_size, layer.row_count - batch_start); + std::fill(link_buffer.begin(), + link_buffer.begin() + current_batch_size * metadata.upper_link_row_bytes, + 0); +#pragma omp parallel for + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); + ++batch_idx) { + const auto row = batch_start + static_cast(batch_idx); + node_buffer[batch_idx] = static_cast(hierarchy.order[start_idx + row]); + auto* link_row = link_buffer.data() + batch_idx * metadata.upper_link_row_bytes; + hnswlib::linklistsizeint list_count = + static_cast(layer.degree); + std::memcpy(link_row, &list_count, sizeof(list_count)); + auto* dst = reinterpret_cast(link_row + sizeof(hnswlib::linklistsizeint)); + if (layer.degree > 0) { + auto* src = host_neighbors.data_handle() + row * layer.degree; + for (size_t j = 0; j < layer.degree; ++j) { + dst[j] = static_cast(hierarchy.order[src[j] + start_idx]); + } + } + } + cuvs::util::write_large_file( + artifact_fd, + node_buffer.data(), + current_batch_size * sizeof(IdxT), + upper_nodes_offset + (layer.node_offset + batch_start) * sizeof(IdxT)); + cuvs::util::write_large_file( + artifact_fd, + link_buffer.data(), + current_batch_size * metadata.upper_link_row_bytes, + upper_links_offset + (layer.link_offset + batch_start) * metadata.upper_link_row_bytes); + } + const auto layer_bytes = layer.row_count * (sizeof(IdxT) + metadata.upper_link_row_bytes); + upper_graph_bytes_written += layer_bytes; + const auto layer_write_elapsed_ms = elapsed_ms_since(layer_write_start_time); + RAFT_LOG_INFO("Layered HNSW: upper layer %zu written in %ld ms (%.2f GiB, %.2f GiB/s)", + layer.level, + layer_write_elapsed_ms, + to_gib(layer_bytes), + throughput_gib_per_s(layer_bytes, layer_write_elapsed_ms)); + }); + const auto upper_layers_elapsed_ms = elapsed_ms_since(upper_layers_start_time); + RAFT_LOG_INFO( + "Layered HNSW: upper layers generated and written in %ld ms (%.2f GiB written, %.2f " + "GiB/s effective)", + upper_layers_elapsed_ms, + to_gib(upper_graph_bytes_written), + throughput_gib_per_s(upper_graph_bytes_written, upper_layers_elapsed_ms)); + } + + const auto total_elapsed_ms = elapsed_ms_since(total_start_time); + RAFT_LOG_INFO( + "Layered HNSW index written to: %s in %ld ms (artifact %.2f GiB, %.2f GiB/s effective)", + artifact_file.string().c_str(), + total_elapsed_ms, + to_gib(final_file_size), + throughput_gib_per_s(final_file_size, total_elapsed_ms)); + return artifact_file.string(); +} + template void serialize_to_hnswlib_from_disk(raft::resources const& res, std::ostream& os_raw, @@ -501,44 +1307,12 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, hnsw_index->get_space(), 1, (graph_degree_int + 1) / 2, params.ef_construction); bool create_hierarchy = params.hierarchy != HnswHierarchy::NONE; - - // create hierarchy order - // sort the points by levels - // roll dice & build histogram - std::vector hist; - std::vector order(n_rows); - std::vector order_bw(n_rows); - std::vector levels(n_rows); - std::vector offsets; - - if (create_hierarchy) { - RAFT_LOG_INFO("Sort points by levels"); - for (int64_t i = 0; i < n_rows; i++) { - auto pt_level = appr_algo->getRandomLevel(appr_algo->mult_); - while (pt_level >= static_cast(hist.size())) - hist.push_back(0); - hist[pt_level]++; - levels[i] = pt_level; - } - - // accumulate - offsets.resize(hist.size() + 1, 0); - for (size_t i = 0; i < hist.size() - 1; i++) { - offsets[i + 1] = offsets[i] + hist[i]; - RAFT_LOG_INFO("Level %zu : %zu", i + 1, size_t(n_rows) - offsets[i + 1]); - } - - // fw/bw indices - for (int64_t i = 0; i < n_rows; i++) { - auto pt_level = levels[i]; - order_bw[i] = offsets[pt_level]; - order[offsets[pt_level]++] = i; - } - } + hnsw_level_plan hierarchy; + if (create_hierarchy) { hierarchy = make_random_hnsw_level_plan(n_rows, *appr_algo, "Level "); } // set last point of the highest level as the entry point - appr_algo->enterpoint_node_ = create_hierarchy ? order.back() : n_rows / 2; - appr_algo->maxlevel_ = create_hierarchy ? hist.size() - 1 : 1; + appr_algo->enterpoint_node_ = create_hierarchy ? hierarchy.order.back() : n_rows / 2; + appr_algo->maxlevel_ = create_hierarchy ? hierarchy.max_level() : 1; // write header information RAFT_LOG_DEBUG("Writing HNSW header: offsetLevel0=%zu, n_rows=%zu, size_data_per_element=%zu", @@ -582,7 +1356,7 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, // host queries auto host_query_set = - raft::make_host_matrix(create_hierarchy ? n_rows - hist[0] : 0, dim); + raft::make_host_matrix(create_hierarchy ? hierarchy.promoted_count() : 0, dim); int64_t d_report_offset = n_rows / 10; // Report progress in 10% steps. int64_t next_report_offset = d_report_offset; @@ -690,11 +1464,11 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, const T* data_row = &dataset_buffer(batch_idx, 0); os.write(reinterpret_cast(data_row), sizeof(T) * dim); - if (create_hierarchy && levels[i] > 0) { + if (create_hierarchy && hierarchy.levels[i] > 0) { // position in query: order_bw[i]-hist[0] - std::copy(data_row, - data_row + dim, - reinterpret_cast(&host_query_set(order_bw[i] - hist[0], 0))); + auto* dst = host_query_set.data_handle() + + (hierarchy.order_bw[i] - hierarchy.hist[0]) * static_cast(dim); + std::copy(data_row, data_row + dim, dst); } // assign original label @@ -732,22 +1506,15 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, // trigger knn builds for all levels std::vector> host_neighbors; if (create_hierarchy) { - for (size_t pt_level = 1; pt_level < hist.size(); pt_level++) { - auto num_pts = n_rows - offsets[pt_level - 1]; - auto neighbor_size = num_pts > appr_algo->M_ ? appr_algo->M_ : num_pts - 1; - host_neighbors.emplace_back(raft::make_host_matrix(num_pts, neighbor_size)); - } - for (size_t pt_level = 1; pt_level < hist.size(); pt_level++) { - RAFT_LOG_INFO("Compute hierarchy neighbors level %zu", pt_level); - auto removed_rows = offsets[pt_level - 1] - offsets[0]; - raft::host_matrix_view sub_query_view( - host_query_set.data_handle() + removed_rows * dim, - host_query_set.extent(0) - removed_rows, - dim); - auto neighbor_view = host_neighbors[pt_level - 1].view(); - all_neighbors_graph( - res, raft::make_const_mdspan(sub_query_view), neighbor_view, index_.metric()); - } + host_neighbors.reserve(hierarchy.hist.size() - 1); + build_hnsw_upper_layer_graphs( + res, + raft::make_host_matrix_view( + host_query_set.data_handle(), host_query_set.extent(0), host_query_set.extent(1)), + hierarchy, + appr_algo->M_, + index_.metric(), + [&](size_t, size_t, auto& neighbors) { host_neighbors.emplace_back(std::move(neighbors)); }); } if (create_hierarchy) { @@ -758,7 +1525,7 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, start_clock = std::chrono::system_clock::now(); for (int64_t i = 0; i < n_rows; i++) { - size_t cur_level = create_hierarchy ? levels[i] : 0; + size_t cur_level = create_hierarchy ? hierarchy.levels[i] : 0; unsigned int linkListSize = create_hierarchy && cur_level > 0 ? appr_algo->size_links_per_element_ * cur_level : 0; os.write(reinterpret_cast(&linkListSize), sizeof(int)); @@ -766,14 +1533,16 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, if (linkListSize) { for (size_t pt_level = 1; pt_level <= cur_level; pt_level++) { auto neighbor_view = host_neighbors[pt_level - 1].view(); - auto my_row = order_bw[i] - offsets[pt_level - 1]; + auto my_row = hierarchy.order_bw[i] - hierarchy.offsets[pt_level - 1]; - IdxT* neighbors = &neighbor_view(my_row, 0); unsigned int extent = neighbor_view.extent(1); os.write(reinterpret_cast(&extent), sizeof(int)); - for (unsigned int j = 0; j < extent; j++) { - const IdxT converted = order[neighbors[j] + offsets[pt_level - 1]]; - os.write(reinterpret_cast(&converted), sizeof(IdxT)); + if (extent > 0) { + IdxT* neighbors = &neighbor_view(my_row, 0); + for (unsigned int j = 0; j < extent; j++) { + const IdxT converted = hierarchy.order[neighbors[j] + hierarchy.offsets[pt_level - 1]]; + os.write(reinterpret_cast(&converted), sizeof(IdxT)); + } } auto remainder = appr_algo->M_ - neighbor_view.extent(1); for (size_t j = 0; j < remainder; j++) { @@ -1088,6 +1857,18 @@ std::unique_ptr> from_cagra( std::filesystem::exists(index_directory) && std::filesystem::is_directory(index_directory), "Directory '%s' does not exist", index_directory.c_str()); + if (params.hierarchy == HnswHierarchy::GPU_LAYERED_ON_DISK) { + RAFT_EXPECTS(dataset.has_value(), + "Layered HNSW serialization requires the original-order dataset."); + auto artifact_path = + serialize_to_layered_hnsw_from_disk(res, params, cagra_index, dataset.value()); + + auto hnsw_index = + std::make_unique>(cagra_index.dim(), cagra_index.metric(), params.hierarchy); + hnsw_index->set_file_descriptor(cuvs::util::file_descriptor(artifact_path, O_RDONLY)); + return hnsw_index; + } + std::string index_filename = (std::filesystem::path(index_directory) / "hnsw_index.bin").string(); @@ -1118,6 +1899,8 @@ std::unique_ptr> from_cagra( return from_cagra(res, params, cagra_index, dataset); } else if (params.hierarchy == HnswHierarchy::GPU) { return from_cagra(res, params, cagra_index, dataset); + } else if (params.hierarchy == HnswHierarchy::GPU_LAYERED_ON_DISK) { + RAFT_FAIL("GPU_LAYERED_ON_DISK requires disk-backed ACE build artifacts."); } else { RAFT_FAIL("Unsupported hierarchy type"); } @@ -1228,6 +2011,13 @@ void serialize(raft::resources const& res, const std::string& filename, const in "Disk-based index file does not exist: %s", source_path.c_str()); + if (idx_impl->hierarchy() == HnswHierarchy::GPU_LAYERED_ON_DISK) { + copy_file_overwrite(source_path, filename); + RAFT_LOG_INFO( + "Copied layered HNSW index from %s to %s", source_path.c_str(), filename.c_str()); + return; + } + // Copy the file to the new location std::filesystem::copy_file( source_path, filename, std::filesystem::copy_options::overwrite_existing); @@ -1241,6 +2031,202 @@ void serialize(raft::resources const& res, const std::string& filename, const in hnswlib_index->saveIndex(filename); } +template +auto deserialize_layered_hnsw(raft::resources const& res, + const index_params& params, + const std::string& artifact_path, + int dim, + cuvs::distance::DistanceType metric) -> std::unique_ptr> +{ + common::nvtx::range fun_scope("hnsw::deserialize_layered"); + cuvs::util::file_descriptor artifact_fd(artifact_path, O_RDONLY); + layered_hnsw_file_header header{}; + cuvs::util::read_large_file(artifact_fd, &header, sizeof(header), 0); + RAFT_EXPECTS(std::strncmp(header.magic, layered_hnsw_magic, sizeof(header.magic)) == 0, + "Invalid layered HNSW artifact magic: %s", + artifact_path.c_str()); + RAFT_EXPECTS(header.version == layered_hnsw_version, + "Unsupported layered HNSW artifact version %u", + header.version); + RAFT_EXPECTS(header.metadata_offset >= sizeof(layered_hnsw_file_header), + "Invalid layered HNSW metadata offset"); + RAFT_EXPECTS(header.metadata_size > 0, "Invalid layered HNSW metadata size"); + + std::string metadata_json(header.metadata_size, '\0'); + cuvs::util::read_large_file( + artifact_fd, metadata_json.data(), metadata_json.size(), header.metadata_offset); + const auto metadata = parse_layered_hnsw_metadata(metadata_json); + + RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); + RAFT_EXPECTS(static_cast(dim) == metadata.dim, + "Layered HNSW artifact dim (%zu) does not match requested dim (%d)", + metadata.dim, + dim); + RAFT_EXPECTS(metadata.layers.size() == static_cast(metadata.maxlevel), + "Layered HNSW artifact has %zu upper layers, expected %d", + metadata.layers.size(), + metadata.maxlevel); + + RAFT_EXPECTS(!params.dataset_path.empty(), + "Layered HNSW deserialization requires index_params.dataset_path"); + + const auto payload_offset = + align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); + const auto levels_offset = payload_offset; + const auto base_links_offset = levels_offset + metadata.levels_bytes; + const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; + const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; + const auto expected_file_size = upper_links_offset + metadata.upper_links_bytes; + const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); + RAFT_EXPECTS(artifact_size >= expected_file_size, + "Layered HNSW artifact is truncated: expected at least %zu bytes, got %zu", + expected_file_size, + artifact_size); + + auto dataset_file = open_layered_dataset_file(params.dataset_path); + + RAFT_EXPECTS(dataset_file.shape.size() == 2 && dataset_file.shape[0] == metadata.n_rows && + dataset_file.shape[1] == metadata.dim, + "Layered HNSW dataset shape mismatch"); + RAFT_EXPECTS(metadata.levels_bytes == metadata.n_rows * sizeof(uint8_t), + "Layered HNSW levels section size mismatch"); + RAFT_EXPECTS(metadata.base_links_bytes == metadata.n_rows * metadata.base_link_row_bytes, + "Layered HNSW base links section size mismatch"); + RAFT_EXPECTS(metadata.upper_nodes_bytes == metadata.upper_nodes_count * sizeof(uint32_t), + "Layered HNSW upper node section size mismatch"); + RAFT_EXPECTS( + metadata.upper_links_bytes == metadata.upper_nodes_count * metadata.upper_link_row_bytes, + "Layered HNSW upper link section size mismatch"); + RAFT_EXPECTS(metadata.base_degree <= metadata.maxM0, + "Layered HNSW base degree (%zu) exceeds maxM0 (%zu)", + metadata.base_degree, + metadata.maxM0); + + std::vector levels_u8(metadata.n_rows); + cuvs::util::read_large_file(artifact_fd, levels_u8.data(), metadata.levels_bytes, levels_offset); + const auto max_level_in_levels = *std::max_element(levels_u8.begin(), levels_u8.end()); + RAFT_EXPECTS(static_cast(max_level_in_levels) == metadata.maxlevel, + "Layered HNSW levels max level (%d) does not match artifact maxlevel (%d)", + static_cast(max_level_in_levels), + metadata.maxlevel); + + auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); + auto appr_algo = std::make_unique::type>>( + hnsw_index->get_space(), metadata.n_rows, metadata.M, metadata.ef_construction); + appr_algo->cur_element_count = metadata.n_rows; + appr_algo->maxlevel_ = metadata.maxlevel; + appr_algo->enterpoint_node_ = metadata.enterpoint_node; + std::fill_n(appr_algo->linkLists_, metadata.n_rows, nullptr); + RAFT_EXPECTS(appr_algo->size_links_level0_ == metadata.base_link_row_bytes, + "Layered HNSW base link row size mismatch"); + RAFT_EXPECTS(appr_algo->size_links_per_element_ == metadata.upper_link_row_bytes, + "Layered HNSW upper link row size mismatch"); + RAFT_EXPECTS(appr_algo->maxM0_ == metadata.maxM0 && appr_algo->maxM_ == metadata.maxM, + "Layered HNSW M parameter mismatch"); + + auto num_threads = + params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; + + RAFT_LOG_INFO("Layered HNSW: initializing in-memory hnswlib index from local dataset"); + const size_t target_batch_bytes = 64 * 1024 * 1024; + const size_t row_bytes = metadata.dim * sizeof(T) + metadata.base_link_row_bytes; + const size_t batch_size = std::max(1, target_batch_bytes / row_bytes); + auto dataset_buffer = + raft::make_host_matrix(static_cast(batch_size), metadata.dim); + std::vector base_link_buffer(batch_size * metadata.base_link_row_bytes); + + for (size_t batch_start = 0; batch_start < metadata.n_rows; batch_start += batch_size) { + const auto current_batch_size = std::min(batch_size, metadata.n_rows - batch_start); + cuvs::util::read_large_file(dataset_file.fd, + dataset_buffer.data_handle(), + current_batch_size * metadata.dim * sizeof(T), + dataset_file.header_size + batch_start * metadata.dim * sizeof(T)); + cuvs::util::read_large_file(artifact_fd, + base_link_buffer.data(), + current_batch_size * metadata.base_link_row_bytes, + base_links_offset + batch_start * metadata.base_link_row_bytes); + + bool link_list_allocation_failed = false; +#pragma omp parallel for num_threads(num_threads) reduction(|| : link_list_allocation_failed) + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { + const auto i = batch_start + static_cast(batch_idx); + auto level = static_cast(levels_u8[i]); + appr_algo->element_levels_[i] = level; + auto ll0 = appr_algo->get_linklist0(i); + memcpy(ll0, + base_link_buffer.data() + batch_idx * metadata.base_link_row_bytes, + metadata.base_link_row_bytes); + memcpy(appr_algo->getDataByInternalId(i), + dataset_buffer.data_handle() + batch_idx * metadata.dim, + appr_algo->data_size_); + *appr_algo->getExternalLabeLp(i) = static_cast(i); + if (level > 0) { + const auto link_list_size = appr_algo->size_links_per_element_ * level + 1; + appr_algo->linkLists_[i] = static_cast(malloc(link_list_size)); + if (appr_algo->linkLists_[i] == nullptr) { + link_list_allocation_failed = true; + continue; + } + memset(appr_algo->linkLists_[i], 0, link_list_size); + } + } + if (link_list_allocation_failed) { + for (size_t i = 0; i < metadata.n_rows; ++i) { + free(appr_algo->linkLists_[i]); + appr_algo->linkLists_[i] = nullptr; + } + throw std::runtime_error("Not enough memory to allocate HNSW upper linklists"); + } + } + + RAFT_LOG_INFO("Layered HNSW: loading upper-layer topology"); + for (const auto& layer : metadata.layers) { + RAFT_LOG_INFO("Layered HNSW: loading layer %zu (%zu rows, degree %zu)", + layer.level, + layer.row_count, + layer.degree); + const auto layer_batch_size = std::min(batch_size, layer.row_count); + std::vector node_buffer(layer_batch_size); + std::vector link_buffer(layer_batch_size * metadata.upper_link_row_bytes); + for (size_t batch_start = 0; batch_start < layer.row_count; batch_start += layer_batch_size) { + const auto current_batch_size = std::min(layer_batch_size, layer.row_count - batch_start); + cuvs::util::read_large_file( + artifact_fd, + node_buffer.data(), + current_batch_size * sizeof(uint32_t), + upper_nodes_offset + (layer.node_offset + batch_start) * sizeof(uint32_t)); + cuvs::util::read_large_file( + artifact_fd, + link_buffer.data(), + current_batch_size * metadata.upper_link_row_bytes, + upper_links_offset + (layer.link_offset + batch_start) * metadata.upper_link_row_bytes); + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); + ++batch_idx) { + const auto node_id = static_cast(node_buffer[batch_idx]); + RAFT_EXPECTS(node_id < metadata.n_rows, + "Invalid upper-layer node id %zu in layered HNSW artifact", + node_id); + RAFT_EXPECTS(layer.level <= static_cast(levels_u8[node_id]), + "Layered HNSW artifact references node %zu at invalid level %zu", + node_id, + layer.level); + } +#pragma omp parallel for num_threads(num_threads) + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); + ++batch_idx) { + const auto node_id = static_cast(node_buffer[batch_idx]); + auto ll = appr_algo->get_linklist(node_id, layer.level); + memcpy(ll, + link_buffer.data() + batch_idx * metadata.upper_link_row_bytes, + metadata.upper_link_row_bytes); + } + } + } + + hnsw_index->set_index(std::move(appr_algo)); + return hnsw_index; +} + template void deserialize(raft::resources const& res, const index_params& params, @@ -1249,6 +2235,12 @@ void deserialize(raft::resources const& res, cuvs::distance::DistanceType metric, index** idx) { + if (params.hierarchy == HnswHierarchy::GPU_LAYERED_ON_DISK) { + auto hnsw_index = deserialize_layered_hnsw(res, params, filename, dim, metric); + *idx = hnsw_index.release(); + return; + } + try { auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); auto appr_algo = std::make_unique::type>>( @@ -1286,6 +2278,13 @@ std::unique_ptr> build(raft::resources const& res, ? std::get(params.graph_build_params) : graph_build_params::ace_params{}; + if (params.hierarchy == HnswHierarchy::GPU_LAYERED_ON_DISK) { + RAFT_EXPECTS(ace_params.use_disk, + "GPU_LAYERED_ON_DISK requires ACE disk mode (ace_params.use_disk = true)"); + RAFT_EXPECTS(!ace_params.build_dir.empty(), + "GPU_LAYERED_ON_DISK requires ace_params.build_dir to be set"); + } + // Create CAGRA index parameters from HNSW parameters cuvs::neighbors::cagra::index_params cagra_params; cagra_params.metric = params.metric; diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 78ae54d9fc..46caef7ae1 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -7,9 +7,16 @@ #include "ann_cagra.cuh" #include +#include +#include +#include #include +#include +#include #include +#include +#include namespace cuvs::neighbors::hnsw { @@ -259,6 +266,238 @@ class AnnHnswAceTest : public ::testing::TestWithParam { std::filesystem::remove_all(temp_dir); } + void testHnswAceLayeredBuildDeserializeSearch() + { + size_t queries_size = ps.n_queries * ps.k; + std::vector indexes_naive(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indexes_naive_dev(queries_size, stream_); + + cuvs::neighbors::naive_knn(handle_, + distances_naive_dev.data(), + indexes_naive_dev.data(), + search_queries.data(), + database_dev.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::update_host(indexes_naive.data(), indexes_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + std::string temp_dir = std::string("/tmp/cuvs_hnsw_ace_layered_test_") + + std::to_string(std::time(nullptr)) + "_" + + std::to_string(reinterpret_cast(this)); + std::filesystem::create_directories(temp_dir); + struct temp_dir_cleanup { + std::string path; + ~temp_dir_cleanup() + { + std::error_code ec; + std::filesystem::remove_all(path, ec); + } + } cleanup{temp_dir}; + + auto database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); + auto queries_host = raft::make_host_matrix(ps.n_queries, ps.dim); + raft::copy(queries_host.data_handle(), search_queries.data(), ps.n_queries * ps.dim, stream_); + raft::resource::sync_stream(handle_); + + const auto dataset_file = (std::filesystem::path(temp_dir) / "dataset.npy").string(); + auto [dataset_fd, dataset_header_size] = cuvs::util::create_numpy_file( + dataset_file, {static_cast(ps.n_rows), static_cast(ps.dim)}); + cuvs::util::write_large_file(dataset_fd, + database_host.data_handle(), + static_cast(ps.n_rows) * ps.dim * sizeof(DataT), + dataset_header_size); + + hnsw::index_params hnsw_params; + hnsw_params.metric = ps.metric; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK; + hnsw_params.M = 32; + hnsw_params.ef_construction = ps.ef_construction; + hnsw_params.dataset_path = dataset_file; + + auto ace_params = graph_build_params::ace_params(); + ace_params.npartitions = ps.npartitions; + ace_params.build_dir = temp_dir; + ace_params.use_disk = true; + ace_params.max_host_memory_gb = ps.max_host_memory_gb; + ace_params.max_gpu_memory_gb = ps.max_gpu_memory_gb; + hnsw_params.graph_build_params = ace_params; + + auto hnsw_index = + hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); + ASSERT_NE(hnsw_index, nullptr); + + const auto artifact_path = hnsw_index->file_path(); + ASSERT_FALSE(artifact_path.empty()); + ASSERT_TRUE(std::filesystem::is_regular_file(artifact_path)); + EXPECT_EQ(std::filesystem::path(artifact_path).filename().string(), "hnsw_index.cuvs"); + size_t cuvs_artifact_count = 0; + for (const auto& entry : std::filesystem::directory_iterator(temp_dir)) { + if (entry.path().extension() == ".cuvs") { ++cuvs_artifact_count; } + } + EXPECT_EQ(cuvs_artifact_count, 1); + EXPECT_FALSE(std::filesystem::exists(std::filesystem::path(temp_dir) / "layered_hnsw")); + + auto indexes_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + auto distances_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + + hnsw::search_params search_params; + search_params.ef = std::max(ps.ef_construction, ps.k * 2); + search_params.num_threads = 1; + + EXPECT_THROW(hnsw::search(handle_, + search_params, + *hnsw_index, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()), + std::exception); + + hnsw::index* deserialized_index = nullptr; + hnsw::deserialize(handle_, hnsw_params, artifact_path, ps.dim, ps.metric, &deserialized_index); + ASSERT_NE(deserialized_index, nullptr); + std::unique_ptr> deserialized_guard(deserialized_index); + + hnsw::search(handle_, + search_params, + *deserialized_guard, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + + std::vector indexes_hnsw_converted(queries_size); + std::vector distances_hnsw(queries_size); + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw_converted[i] = static_cast(indexes_hnsw_host.data_handle()[i]); + distances_hnsw[i] = distances_hnsw_host.data_handle()[i]; + } + + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_naive, + indexes_hnsw_converted, + distances_naive, + distances_hnsw, + ps.n_queries, + ps.k, + 0.003, + ps.min_recall)) + << "Layered HNSW deserialize and search failed recall check"; + + const auto copied_artifact = + (std::filesystem::path(temp_dir) / "copied_layered" / "hnsw_index.cuvs").string(); + hnsw::serialize(handle_, copied_artifact, *hnsw_index); + EXPECT_TRUE(std::filesystem::is_regular_file(copied_artifact)); + size_t copied_file_count = 0; + for (const auto& entry : std::filesystem::directory_iterator( + std::filesystem::path(copied_artifact).parent_path())) { + if (entry.is_regular_file()) { ++copied_file_count; } + } + EXPECT_EQ(copied_file_count, 1); + + hnsw::index* copied_index = nullptr; + hnsw::deserialize(handle_, hnsw_params, copied_artifact, ps.dim, ps.metric, &copied_index); + ASSERT_NE(copied_index, nullptr); + std::unique_ptr> copied_guard(copied_index); + + hnsw::search(handle_, + search_params, + *copied_guard, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw_converted[i] = static_cast(indexes_hnsw_host.data_handle()[i]); + distances_hnsw[i] = distances_hnsw_host.data_handle()[i]; + } + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_naive, + indexes_hnsw_converted, + distances_naive, + distances_hnsw, + ps.n_queries, + ps.k, + 0.003, + ps.min_recall)) + << "Copied layered HNSW artifact deserialize and search failed recall check"; + + const auto bad_artifact = + (std::filesystem::path(temp_dir) / "bad_layered" / "bad_magic.cuvs").string(); + std::filesystem::create_directories(std::filesystem::path(bad_artifact).parent_path()); + { + std::ofstream bad_file(bad_artifact, std::ios::binary); + std::array bad_bytes{}; + bad_file.write(bad_bytes.data(), bad_bytes.size()); + } + hnsw::index* bad_index = nullptr; + EXPECT_THROW( + hnsw::deserialize(handle_, hnsw_params, bad_artifact, ps.dim, ps.metric, &bad_index), + std::exception); + + const auto bad_version_artifact = + (std::filesystem::path(temp_dir) / "bad_layered" / "bad_version.cuvs").string(); + std::filesystem::copy_file( + copied_artifact, bad_version_artifact, std::filesystem::copy_options::overwrite_existing); + constexpr std::streamoff layered_header_version_offset = 32; + { + std::fstream bad_version_file(bad_version_artifact, + std::ios::in | std::ios::out | std::ios::binary); + const uint32_t bad_version = 999; + bad_version_file.seekp(layered_header_version_offset); + bad_version_file.write(reinterpret_cast(&bad_version), sizeof(bad_version)); + } + hnsw::index* bad_version_index = nullptr; + EXPECT_THROW( + hnsw::deserialize( + handle_, hnsw_params, bad_version_artifact, ps.dim, ps.metric, &bad_version_index), + std::exception); + + auto missing_dataset_params = hnsw_params; + missing_dataset_params.dataset_path.clear(); + hnsw::index* missing_dataset_index = nullptr; + EXPECT_THROW(hnsw::deserialize(handle_, + missing_dataset_params, + copied_artifact, + ps.dim, + ps.metric, + &missing_dataset_index), + std::exception); + + const auto truncated_artifact = + (std::filesystem::path(temp_dir) / "bad_layered" / "truncated.cuvs").string(); + std::filesystem::copy_file( + copied_artifact, truncated_artifact, std::filesystem::copy_options::overwrite_existing); + std::filesystem::resize_file(truncated_artifact, 128); + hnsw::index* truncated_index = nullptr; + EXPECT_THROW(hnsw::deserialize( + handle_, hnsw_params, truncated_artifact, ps.dim, ps.metric, &truncated_index), + std::exception); + + const auto wrong_dataset_file = + (std::filesystem::path(temp_dir) / "bad_layered" / "wrong_dataset.npy").string(); + auto [wrong_dataset_fd, wrong_dataset_header_size] = cuvs::util::create_numpy_file( + wrong_dataset_file, {static_cast(ps.n_rows - 1), static_cast(ps.dim)}); + cuvs::util::write_large_file(wrong_dataset_fd, + database_host.data_handle(), + static_cast(ps.n_rows - 1) * ps.dim * sizeof(DataT), + wrong_dataset_header_size); + auto wrong_dataset_params = hnsw_params; + wrong_dataset_params.dataset_path = wrong_dataset_file; + hnsw::index* wrong_dataset_index = nullptr; + EXPECT_THROW( + hnsw::deserialize( + handle_, wrong_dataset_params, copied_artifact, ps.dim, ps.metric, &wrong_dataset_index), + std::exception); + } + void SetUp() override { database_dev.resize(((size_t)ps.n_rows) * ps.dim, stream_); @@ -323,8 +562,26 @@ inline std::vector generate_hnsw_ace_memory_fallback_inputs() }; } +inline std::vector generate_hnsw_ace_layered_inputs() +{ + return { + {10, // n_queries + 5000, // n_rows + 64, // dim + 10, // k + 2, // npartitions + 100, // ef_construction + true, // use_disk + cuvs::distance::DistanceType::L2Expanded, + 0.9, // min_recall + 0.0, // max_host_memory_gb + 0.0} // max_gpu_memory_gb + }; +} + const std::vector hnsw_ace_inputs = generate_hnsw_ace_inputs(); const std::vector hnsw_ace_memory_fallback_inputs = generate_hnsw_ace_memory_fallback_inputs(); +const std::vector hnsw_ace_layered_inputs = generate_hnsw_ace_layered_inputs(); } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu index 57b16d74e9..97a6255300 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -23,4 +23,14 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceMemoryFallbackTest, AnnHnswAceMemoryFallbackTest_float, ::testing::ValuesIn(hnsw_ace_memory_fallback_inputs)); +typedef AnnHnswAceTest AnnHnswAceLayeredTest_float; +TEST_P(AnnHnswAceLayeredTest_float, AnnHnswAceLayeredBuildDeserializeSearch) +{ + this->testHnswAceLayeredBuildDeserializeSearch(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, + AnnHnswAceLayeredTest_float, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu index 38f75b2afa..c77391ba28 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -12,4 +12,24 @@ TEST_P(AnnHnswAceTest_half, AnnHnswAceBuild) { this->testHnswAceBuild(); } INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_half, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_half; +TEST_P(AnnHnswAceMemoryFallbackTest_half, AnnHnswAceMemoryLimitFallback) +{ + this->testHnswAceMemoryLimitFallback(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMemoryFallbackTest, + AnnHnswAceMemoryFallbackTest_half, + ::testing::ValuesIn(hnsw_ace_memory_fallback_inputs)); + +typedef AnnHnswAceTest AnnHnswAceLayeredTest_half; +TEST_P(AnnHnswAceLayeredTest_half, AnnHnswAceLayeredBuildDeserializeSearch) +{ + this->testHnswAceLayeredBuildDeserializeSearch(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, + AnnHnswAceLayeredTest_half, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu index 279df6555f..4ee6ecc5e5 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -14,4 +14,24 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_int8_t, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_int8_t; +TEST_P(AnnHnswAceMemoryFallbackTest_int8_t, AnnHnswAceMemoryLimitFallback) +{ + this->testHnswAceMemoryLimitFallback(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMemoryFallbackTest, + AnnHnswAceMemoryFallbackTest_int8_t, + ::testing::ValuesIn(hnsw_ace_memory_fallback_inputs)); + +typedef AnnHnswAceTest AnnHnswAceLayeredTest_int8_t; +TEST_P(AnnHnswAceLayeredTest_int8_t, AnnHnswAceLayeredBuildDeserializeSearch) +{ + this->testHnswAceLayeredBuildDeserializeSearch(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, + AnnHnswAceLayeredTest_int8_t, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu index 7e68dc4b17..1ef791df8e 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -14,4 +14,24 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_uint8_t, ::testing::ValuesIn(hnsw_ace_inputs)); +typedef AnnHnswAceTest AnnHnswAceMemoryFallbackTest_uint8_t; +TEST_P(AnnHnswAceMemoryFallbackTest_uint8_t, AnnHnswAceMemoryLimitFallback) +{ + this->testHnswAceMemoryLimitFallback(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMemoryFallbackTest, + AnnHnswAceMemoryFallbackTest_uint8_t, + ::testing::ValuesIn(hnsw_ace_memory_fallback_inputs)); + +typedef AnnHnswAceTest AnnHnswAceLayeredTest_uint8_t; +TEST_P(AnnHnswAceLayeredTest_uint8_t, AnnHnswAceLayeredBuildDeserializeSearch) +{ + this->testHnswAceLayeredBuildDeserializeSearch(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, + AnnHnswAceLayeredTest_uint8_t, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 034b0b3d96..930f898266 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -35,6 +35,7 @@ add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) add_executable(HNSW_ACE_EXAMPLE src/hnsw_ace_example.cu) +add_executable(HNSW_ACE_LAYERED_EXAMPLE src/hnsw_ace_layered_example.cu) add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) add_executable(VAMANA_EXAMPLE src/vamana_example.cu) @@ -52,6 +53,7 @@ target_link_libraries( DYNAMIC_BATCHING_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads ) target_link_libraries(HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(HNSW_ACE_LAYERED_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(IVF_PQ_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(IVF_FLAT_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(VAMANA_EXAMPLE PRIVATE cuvs::cuvs $) diff --git a/examples/cpp/src/hnsw_ace_layered_example.cu b/examples/cpp/src/hnsw_ace_layered_example.cu new file mode 100644 index 0000000000..02f735ef37 --- /dev/null +++ b/examples/cpp/src/hnsw_ace_layered_example.cu @@ -0,0 +1,278 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Large graphs exceeding the memory capacity can be built using the Augmented Core Extraction (ACE) +// algorithm, which partitions the dataset. The resulting HNSW index is too large to fit in memory +// as well. Thus, the index needs to be transferred to a search server with enough memory. +// +// HNSWHierarchy::GPU_LAYERED_ON_DISK is a special hierarchy that builds a layered HNSW index on +// disk. It emits one topology-only artifact, hnsw_index.cuvs. The dataset remains separate and does +// not need to be transferred to the search server, which typically has the dataset locally. +// +// This example demonstrates how to build a layered HNSW index with ACE: +// +// 1. Optionally quantize the dataset and queries to int8. +// 2. Build a single-file layered HNSW artifact with ACE using hnsw::build. +// 3. Deserialize the layered HNSW artifact using hnsw::deserialize. +// 4. Search the in-memory HNSW index. +// +// Layered-on-disk layout: +// +// index_dir/hnsw_index.cuvs +// fixed header + metadata JSON +// levels: uint8 [N], max HNSW level for each original row id +// base links: hnswlib-ready [count + uint32 neighbors padded to maxM0] +// upper nodes + upper links: hnswlib-ready upper-layer topology +// +// The transferred index artifact is topology-only. The dataset is loaded locally during +// deserialization from hnsw::index_params::dataset_path. The loader supports .npy and ANN benchmark +// *.bin datasets; this example writes a local dataset .npy only to make the demo self-contained. +// +// Layer 0 rows and neighbor IDs are original dataset row IDs. Upper layers are generated with the +// same level/order/KNN logic as serialize_to_hnswlib_from_disk, then stored as hnswlib-ready link +// rows so deserialization does no graph remapping or link padding on the search node. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include "common.cuh" + +// When 1, scalar-quantize the float dataset to int8. +#define HNSW_ACE_LAYERED_USE_QUANTIZATION 1 + +namespace { + +constexpr const char* kBuildDir = "/tmp/hnsw_ace_layered"; + +template +std::string write_local_dataset(raft::host_matrix_view dataset, + const std::string& path) +{ + auto [fd, header_size] = cuvs::util::create_numpy_file( + path, {static_cast(dataset.extent(0)), static_cast(dataset.extent(1))}); + cuvs::util::write_large_file( + fd, dataset.data_handle(), dataset.extent(0) * dataset.extent(1) * sizeof(T), header_size); + return path; +} + +template +struct quantized_pair { + raft::host_matrix dataset; + raft::host_matrix queries; +}; + +quantized_pair quantize_dataset(raft::device_resources const& dev_resources, + raft::host_matrix_view dataset_float, + raft::host_matrix_view queries_float) +{ + std::cout << " quantize_dataset: training scalar quantizer (float -> int8)" << std::endl; + cuvs::preprocessing::quantize::scalar::params qp; + auto quantizer = cuvs::preprocessing::quantize::scalar::train(dev_resources, qp, dataset_float); + + auto dataset_i8 = + raft::make_host_matrix(dataset_float.extent(0), dataset_float.extent(1)); + cuvs::preprocessing::quantize::scalar::transform( + dev_resources, quantizer, dataset_float, dataset_i8.view()); + + auto queries_i8 = + raft::make_host_matrix(queries_float.extent(0), queries_float.extent(1)); + cuvs::preprocessing::quantize::scalar::transform( + dev_resources, quantizer, queries_float, queries_i8.view()); + + return {std::move(dataset_i8), std::move(queries_i8)}; +} + +auto make_hnsw_ace_params(const std::string& build_dir, const std::string& dataset_path) + -> cuvs::neighbors::hnsw::index_params +{ + using namespace cuvs::neighbors; + + hnsw::index_params hnsw_params; + hnsw_params.metric = cuvs::distance::DistanceType::L2Expanded; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK; + hnsw_params.M = 32; + hnsw_params.dataset_path = dataset_path; // Override this path on the search server. + hnsw_params.ef_construction = 120; + + auto ace_params = hnsw::graph_build_params::ace_params(); + ace_params.npartitions = 4; + ace_params.use_disk = true; + ace_params.build_dir = build_dir; + hnsw_params.graph_build_params = ace_params; + + return hnsw_params; +} + +template +auto hnsw_build(raft::device_resources const& dev_resources, + const cuvs::neighbors::hnsw::index_params& hnsw_params, + raft::host_matrix_view dataset) -> std::string +{ + using namespace cuvs::neighbors; + + auto hnsw_index = hnsw::build(dev_resources, hnsw_params, dataset); + const auto artifact_path = hnsw_index->file_path(); + if (artifact_path.empty()) { + throw std::runtime_error("Expected layered HNSW build to return an artifact path."); + } + std::cout << " hnsw_build: layered artifact written to " << artifact_path << std::endl; + return artifact_path; +} + +template +auto hnsw_deserialize(raft::device_resources const& dev_resources, + const cuvs::neighbors::hnsw::index_params& hnsw_params, + const std::string& artifact_path, + int64_t dim) -> std::unique_ptr> +{ + using namespace cuvs::neighbors; + + hnsw::index* deserialized_index = nullptr; + // Set params.dataset_path to the local dataset path to load the dataset from disk. + // hnsw_params.dataset_path = "/tmp/dataset.npy"; + hnsw::deserialize( + dev_resources, hnsw_params, artifact_path, dim, hnsw_params.metric, &deserialized_index); + return std::unique_ptr>(deserialized_index); +} + +template +void hnsw_search(raft::device_resources const& dev_resources, + const cuvs::neighbors::hnsw::index& hnsw_index, + raft::host_matrix_view queries, + int64_t topk = 12) +{ + using namespace cuvs::neighbors; + + const int64_t n_queries = queries.extent(0); + auto indices_hnsw_host = raft::make_host_matrix(n_queries, topk); + auto distances_hnsw_host = raft::make_host_matrix(n_queries, topk); + + hnsw::search_params search_params; + search_params.ef = std::max(200, static_cast(topk) * 2); + search_params.num_threads = 1; + + hnsw::search(dev_resources, + search_params, + hnsw_index, + queries, + indices_hnsw_host.view(), + distances_hnsw_host.view()); + + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + auto neighbors_host = raft::make_host_matrix(n_queries, topk); + for (int64_t i = 0; i < n_queries; ++i) { + for (int64_t j = 0; j < topk; ++j) { + neighbors_host(i, j) = static_cast(indices_hnsw_host(i, j)); + } + } + + raft::copy(neighbors.data_handle(), + neighbors_host.data_handle(), + n_queries * topk, + raft::resource::get_cuda_stream(dev_resources)); + raft::copy(distances.data_handle(), + distances_hnsw_host.data_handle(), + n_queries * topk, + raft::resource::get_cuda_stream(dev_resources)); + raft::resource::sync_stream(dev_resources); + + print_results(dev_resources, neighbors.view(), distances.view()); +} + +} // namespace + +int main() +{ + raft::device_resources dev_resources; + + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), + 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(pool_mr); + +#if HNSW_ACE_LAYERED_USE_QUANTIZATION + std::cout << "[stage 1] Generate and quantize dataset (float -> int8)" << std::endl; +#else + std::cout << "[stage 1] Generate dataset (float)" << std::endl; +#endif + + int64_t n_samples = 10000; + int64_t n_dim = 90; + int64_t n_queries = 10; + auto dataset = raft::make_device_matrix(dev_resources, n_samples, n_dim); + auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); + generate_dataset(dev_resources, dataset.view(), queries.view()); + + auto dataset_host = raft::make_host_matrix(n_samples, n_dim); + auto queries_host = raft::make_host_matrix(n_queries, n_dim); + raft::copy(dataset_host.data_handle(), + dataset.data_handle(), + dataset.extent(0) * dataset.extent(1), + raft::resource::get_cuda_stream(dev_resources)); + raft::copy(queries_host.data_handle(), + queries.data_handle(), + queries.extent(0) * queries.extent(1), + raft::resource::get_cuda_stream(dev_resources)); + raft::resource::sync_stream(dev_resources); + + auto dataset_host_view = raft::make_host_matrix_view( + dataset_host.data_handle(), n_samples, n_dim); + auto queries_host_view = raft::make_host_matrix_view( + queries_host.data_handle(), n_queries, n_dim); + + std::filesystem::create_directories(kBuildDir); + +#if HNSW_ACE_LAYERED_USE_QUANTIZATION + auto q = quantize_dataset(dev_resources, dataset_host_view, queries_host_view); + auto dataset_i8_view = raft::make_host_matrix_view( + q.dataset.data_handle(), n_samples, n_dim); + auto queries_i8_view = raft::make_host_matrix_view( + q.queries.data_handle(), n_queries, n_dim); + auto dataset_path = write_local_dataset(dataset_i8_view, std::string{kBuildDir} + "/dataset.npy"); + auto hnsw_params = make_hnsw_ace_params(kBuildDir, dataset_path); + + std::cout << "[stage 2] Build layered HNSW index with ACE" << std::endl; + auto artifact_path = hnsw_build(dev_resources, hnsw_params, dataset_i8_view); + + std::cout << "[stage 3] Deserialize layered HNSW index" << std::endl; + auto hnsw_index = hnsw_deserialize(dev_resources, hnsw_params, artifact_path, n_dim); + + std::cout << "[stage 4] Search HNSW index" << std::endl; + hnsw_search(dev_resources, *hnsw_index, queries_i8_view); +#else + auto dataset_path = + write_local_dataset(dataset_host_view, std::string{kBuildDir} + "/dataset.npy"); + auto hnsw_params = make_hnsw_ace_params(kBuildDir, dataset_path); + + std::cout << "[stage 2] Build layered HNSW index with ACE" << std::endl; + auto artifact_path = hnsw_build(dev_resources, hnsw_params, dataset_host_view); + + std::cout << "[stage 3] Deserialize layered HNSW index" << std::endl; + auto hnsw_index = hnsw_deserialize(dev_resources, hnsw_params, artifact_path, n_dim); + + std::cout << "[stage 4] Search HNSW index" << std::endl; + hnsw_search(dev_resources, *hnsw_index, queries_host_view); +#endif + + return 0; +} From a11d9ee6ae0ee1d00f666fc90b40df2de09b64e8 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 27 May 2026 10:01:25 +0200 Subject: [PATCH 02/12] Improve deserialization logging --- cpp/src/neighbors/detail/hnsw.hpp | 101 ++++++++++++++++++++++++++++-- 1 file changed, 97 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 44dc4e688e..c5b86076c9 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -389,6 +389,11 @@ inline auto elapsed_ms_since(std::chrono::steady_clock::time_point start) -> int .count(); } +inline auto elapsed_ms(std::chrono::steady_clock::duration elapsed) -> double +{ + return std::chrono::duration(elapsed).count(); +} + inline auto to_gib(size_t bytes) -> double { return static_cast(bytes) / (1024.0 * 1024.0 * 1024.0); @@ -399,6 +404,13 @@ inline auto throughput_gib_per_s(size_t bytes, int64_t elapsed_ms) -> double return elapsed_ms > 0 ? to_gib(bytes) / (elapsed_ms / 1000.0) : 0.0; } +inline auto throughput_gib_per_s(size_t bytes, std::chrono::steady_clock::duration elapsed) + -> double +{ + const auto elapsed_s = std::chrono::duration(elapsed).count(); + return elapsed_s > 0.0 ? to_gib(bytes) / elapsed_s : 0.0; +} + inline auto make_hnsw_level_plan_from_levels(size_t n_rows, std::vector&& levels, bool build_reverse_order, @@ -2039,6 +2051,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, cuvs::distance::DistanceType metric) -> std::unique_ptr> { common::nvtx::range fun_scope("hnsw::deserialize_layered"); + const auto total_start_time = std::chrono::steady_clock::now(); + const auto metadata_start_time = std::chrono::steady_clock::now(); cuvs::util::file_descriptor artifact_fd(artifact_path, O_RDONLY); layered_hnsw_file_header header{}; cuvs::util::read_large_file(artifact_fd, &header, sizeof(header), 0); @@ -2055,7 +2069,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, std::string metadata_json(header.metadata_size, '\0'); cuvs::util::read_large_file( artifact_fd, metadata_json.data(), metadata_json.size(), header.metadata_offset); - const auto metadata = parse_layered_hnsw_metadata(metadata_json); + const auto metadata = parse_layered_hnsw_metadata(metadata_json); + const auto metadata_elapsed_ms = elapsed_ms_since(metadata_start_time); RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); RAFT_EXPECTS(static_cast(dim) == metadata.dim, @@ -2082,8 +2097,15 @@ auto deserialize_layered_hnsw(raft::resources const& res, "Layered HNSW artifact is truncated: expected at least %zu bytes, got %zu", expected_file_size, artifact_size); + RAFT_LOG_INFO("Layered HNSW load: metadata read in %ld ms (rows=%zu dim=%zu artifact=%.2f GiB)", + metadata_elapsed_ms, + metadata.n_rows, + metadata.dim, + to_gib(artifact_size)); - auto dataset_file = open_layered_dataset_file(params.dataset_path); + const auto dataset_open_start_time = std::chrono::steady_clock::now(); + auto dataset_file = open_layered_dataset_file(params.dataset_path); + const auto dataset_open_elapsed_ms = elapsed_ms_since(dataset_open_start_time); RAFT_EXPECTS(dataset_file.shape.size() == 2 && dataset_file.shape[0] == metadata.n_rows && dataset_file.shape[1] == metadata.dim, @@ -2102,6 +2124,11 @@ auto deserialize_layered_hnsw(raft::resources const& res, metadata.base_degree, metadata.maxM0); + RAFT_LOG_INFO("Layered HNSW load: dataset header validated in %ld ms (%s)", + dataset_open_elapsed_ms, + params.dataset_path.c_str()); + + const auto levels_start_time = std::chrono::steady_clock::now(); std::vector levels_u8(metadata.n_rows); cuvs::util::read_large_file(artifact_fd, levels_u8.data(), metadata.levels_bytes, levels_offset); const auto max_level_in_levels = *std::max_element(levels_u8.begin(), levels_u8.end()); @@ -2109,9 +2136,14 @@ auto deserialize_layered_hnsw(raft::resources const& res, "Layered HNSW levels max level (%d) does not match artifact maxlevel (%d)", static_cast(max_level_in_levels), metadata.maxlevel); + const auto levels_elapsed_ms = elapsed_ms_since(levels_start_time); + RAFT_LOG_INFO("Layered HNSW load: levels read in %ld ms (%.2f MiB)", + levels_elapsed_ms, + static_cast(metadata.levels_bytes) / (1024.0 * 1024.0)); - auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); - auto appr_algo = std::make_unique::type>>( + const auto allocation_start_time = std::chrono::steady_clock::now(); + auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); + auto appr_algo = std::make_unique::type>>( hnsw_index->get_space(), metadata.n_rows, metadata.M, metadata.ef_construction); appr_algo->cur_element_count = metadata.n_rows; appr_algo->maxlevel_ = metadata.maxlevel; @@ -2123,6 +2155,12 @@ auto deserialize_layered_hnsw(raft::resources const& res, "Layered HNSW upper link row size mismatch"); RAFT_EXPECTS(appr_algo->maxM0_ == metadata.maxM0 && appr_algo->maxM_ == metadata.maxM, "Layered HNSW M parameter mismatch"); + const auto allocation_elapsed_ms = elapsed_ms_since(allocation_start_time); + RAFT_LOG_INFO( + "Layered HNSW load: allocated hnswlib storage in %ld ms (base=%.2f GiB upper=%.2f GiB)", + allocation_elapsed_ms, + to_gib(metadata.n_rows * appr_algo->size_data_per_element_), + to_gib(metadata.upper_nodes_count * metadata.upper_link_row_bytes)); auto num_threads = params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; @@ -2135,18 +2173,31 @@ auto deserialize_layered_hnsw(raft::resources const& res, raft::make_host_matrix(static_cast(batch_size), metadata.dim); std::vector base_link_buffer(batch_size * metadata.base_link_row_bytes); + const auto base_start_time = std::chrono::steady_clock::now(); + std::chrono::steady_clock::duration dataset_read_time{}; + std::chrono::steady_clock::duration base_link_read_time{}; + std::chrono::steady_clock::duration base_copy_time{}; + size_t dataset_bytes_read = 0; + size_t base_link_bytes_read = 0; for (size_t batch_start = 0; batch_start < metadata.n_rows; batch_start += batch_size) { const auto current_batch_size = std::min(batch_size, metadata.n_rows - batch_start); + auto batch_timer = std::chrono::steady_clock::now(); cuvs::util::read_large_file(dataset_file.fd, dataset_buffer.data_handle(), current_batch_size * metadata.dim * sizeof(T), dataset_file.header_size + batch_start * metadata.dim * sizeof(T)); + dataset_read_time += std::chrono::steady_clock::now() - batch_timer; + dataset_bytes_read += current_batch_size * metadata.dim * sizeof(T); + batch_timer = std::chrono::steady_clock::now(); cuvs::util::read_large_file(artifact_fd, base_link_buffer.data(), current_batch_size * metadata.base_link_row_bytes, base_links_offset + batch_start * metadata.base_link_row_bytes); + base_link_read_time += std::chrono::steady_clock::now() - batch_timer; + base_link_bytes_read += current_batch_size * metadata.base_link_row_bytes; bool link_list_allocation_failed = false; + batch_timer = std::chrono::steady_clock::now(); #pragma omp parallel for num_threads(num_threads) reduction(|| : link_list_allocation_failed) for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { const auto i = batch_start + static_cast(batch_idx); @@ -2177,10 +2228,30 @@ auto deserialize_layered_hnsw(raft::resources const& res, } throw std::runtime_error("Not enough memory to allocate HNSW upper linklists"); } + base_copy_time += std::chrono::steady_clock::now() - batch_timer; } + const auto base_elapsed_ms = elapsed_ms_since(base_start_time); + RAFT_LOG_INFO( + "Layered HNSW load: base layer initialized in %ld ms " + "(dataset read %.2f ms %.2f GiB %.2f GiB/s, links read %.2f ms %.2f GiB %.2f GiB/s, " + "copy %.2f ms)", + base_elapsed_ms, + elapsed_ms(dataset_read_time), + to_gib(dataset_bytes_read), + throughput_gib_per_s(dataset_bytes_read, dataset_read_time), + elapsed_ms(base_link_read_time), + to_gib(base_link_bytes_read), + throughput_gib_per_s(base_link_bytes_read, base_link_read_time), + elapsed_ms(base_copy_time)); RAFT_LOG_INFO("Layered HNSW: loading upper-layer topology"); + const auto upper_start_time = std::chrono::steady_clock::now(); + std::chrono::steady_clock::duration upper_read_time{}; + std::chrono::steady_clock::duration upper_validate_time{}; + std::chrono::steady_clock::duration upper_copy_time{}; + size_t upper_bytes_read = 0; for (const auto& layer : metadata.layers) { + const auto layer_start_time = std::chrono::steady_clock::now(); RAFT_LOG_INFO("Layered HNSW: loading layer %zu (%zu rows, degree %zu)", layer.level, layer.row_count, @@ -2190,6 +2261,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, std::vector link_buffer(layer_batch_size * metadata.upper_link_row_bytes); for (size_t batch_start = 0; batch_start < layer.row_count; batch_start += layer_batch_size) { const auto current_batch_size = std::min(layer_batch_size, layer.row_count - batch_start); + auto batch_timer = std::chrono::steady_clock::now(); cuvs::util::read_large_file( artifact_fd, node_buffer.data(), @@ -2200,6 +2272,9 @@ auto deserialize_layered_hnsw(raft::resources const& res, link_buffer.data(), current_batch_size * metadata.upper_link_row_bytes, upper_links_offset + (layer.link_offset + batch_start) * metadata.upper_link_row_bytes); + upper_read_time += std::chrono::steady_clock::now() - batch_timer; + upper_bytes_read += current_batch_size * (sizeof(uint32_t) + metadata.upper_link_row_bytes); + batch_timer = std::chrono::steady_clock::now(); for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { const auto node_id = static_cast(node_buffer[batch_idx]); @@ -2211,6 +2286,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, node_id, layer.level); } + upper_validate_time += std::chrono::steady_clock::now() - batch_timer; + batch_timer = std::chrono::steady_clock::now(); #pragma omp parallel for num_threads(num_threads) for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { @@ -2220,10 +2297,26 @@ auto deserialize_layered_hnsw(raft::resources const& res, link_buffer.data() + batch_idx * metadata.upper_link_row_bytes, metadata.upper_link_row_bytes); } + upper_copy_time += std::chrono::steady_clock::now() - batch_timer; } + RAFT_LOG_INFO("Layered HNSW load: layer %zu loaded in %ld ms", + layer.level, + elapsed_ms_since(layer_start_time)); } + const auto upper_elapsed_ms = elapsed_ms_since(upper_start_time); + RAFT_LOG_INFO( + "Layered HNSW load: upper layers loaded in %ld ms " + "(read %.2f ms %.2f GiB %.2f GiB/s, validate %.2f ms, copy %.2f ms)", + upper_elapsed_ms, + elapsed_ms(upper_read_time), + to_gib(upper_bytes_read), + throughput_gib_per_s(upper_bytes_read, upper_read_time), + elapsed_ms(upper_validate_time), + elapsed_ms(upper_copy_time)); hnsw_index->set_index(std::move(appr_algo)); + RAFT_LOG_INFO("Layered HNSW load: total deserialize completed in %ld ms", + elapsed_ms_since(total_start_time)); return hnsw_index; } From a95b0e0c0014982f469e62e757a349798eaedb2b Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 27 May 2026 11:58:21 +0200 Subject: [PATCH 03/12] Use ace prefix in benchmarking consistently --- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu index b535ced763..d4791dde11 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu @@ -64,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 (conf.contains("use_disk")) { ace_params.use_disk = conf.at("use_disk"); } + 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 (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. From 47e2d858a410030913bd587ed2c11661694ad95d Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 27 May 2026 11:59:10 +0200 Subject: [PATCH 04/12] Validate metadata before allocating --- cpp/src/neighbors/detail/hnsw.hpp | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index c5b86076c9..d048d322a2 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -2065,6 +2065,14 @@ auto deserialize_layered_hnsw(raft::resources const& res, RAFT_EXPECTS(header.metadata_offset >= sizeof(layered_hnsw_file_header), "Invalid layered HNSW metadata offset"); RAFT_EXPECTS(header.metadata_size > 0, "Invalid layered HNSW metadata size"); + const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); + RAFT_EXPECTS(header.metadata_offset <= artifact_size && + header.metadata_size <= artifact_size - header.metadata_offset, + "Layered HNSW metadata range is outside artifact: offset=%zu size=%zu " + "artifact=%zu", + static_cast(header.metadata_offset), + static_cast(header.metadata_size), + artifact_size); std::string metadata_json(header.metadata_size, '\0'); cuvs::util::read_large_file( @@ -2092,7 +2100,6 @@ auto deserialize_layered_hnsw(raft::resources const& res, const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; const auto expected_file_size = upper_links_offset + metadata.upper_links_bytes; - const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); RAFT_EXPECTS(artifact_size >= expected_file_size, "Layered HNSW artifact is truncated: expected at least %zu bytes, got %zu", expected_file_size, From 21ce339abadb11c2bcbeca3c389614877f4e72d1 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Mon, 1 Jun 2026 10:12:45 +0200 Subject: [PATCH 05/12] Store layered base topology by original node ID - Add base node IDs for sequential access. - Scattered writes happen only in deserialization step using host memory. --- cpp/src/neighbors/detail/hnsw.hpp | 351 ++++++++++--------- examples/cpp/src/hnsw_ace_layered_example.cu | 8 +- 2 files changed, 185 insertions(+), 174 deletions(-) diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index d048d322a2..d090de07f8 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -40,8 +40,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -519,6 +519,7 @@ struct layered_hnsw_file_metadata { size_t ef_construction = 0; size_t base_degree = 0; size_t levels_bytes = 0; + size_t base_nodes_bytes = 0; size_t base_link_row_bytes = 0; size_t base_links_bytes = 0; size_t upper_nodes_count = 0; @@ -580,12 +581,13 @@ auto make_layered_hnsw_metadata_json(const layered_hnsw_file_metadata& metadata, std::stringstream os; os << "{\n"; os << " \"format\": \"cuvs_hnsw_layered\",\n"; - os << " \"version\": 1,\n"; + os << " \"version\": " << layered_hnsw_version << ",\n"; os << " \"n_rows\": " << metadata.n_rows << ",\n"; os << " \"dim\": " << metadata.dim << ",\n"; os << " \"dtype\": \"" << layered_dtype_name() << "\",\n"; os << " \"metric\": \"" << metric_name(metric) << "\",\n"; os << " \"row_order\": \"original_id\",\n"; + os << " \"base_link_order\": \"base_nodes\",\n"; os << " \"M\": " << metadata.M << ",\n"; os << " \"maxM\": " << metadata.maxM << ",\n"; os << " \"maxM0\": " << metadata.maxM0 << ",\n"; @@ -595,6 +597,7 @@ auto make_layered_hnsw_metadata_json(const layered_hnsw_file_metadata& metadata, os << " \"enterpoint_node\": " << metadata.enterpoint_node << ",\n"; os << " \"base_degree\": " << metadata.base_degree << ",\n"; os << " \"levels_bytes\": " << metadata.levels_bytes << ",\n"; + os << " \"base_nodes_bytes\": " << metadata.base_nodes_bytes << ",\n"; os << " \"base_link_row_bytes\": " << metadata.base_link_row_bytes << ",\n"; os << " \"base_links_bytes\": " << metadata.base_links_bytes << ",\n"; os << " \"upper_nodes_count\": " << metadata.upper_nodes_count << ",\n"; @@ -665,6 +668,7 @@ inline auto parse_layered_hnsw_metadata(const std::string& json) -> layered_hnsw metadata.enterpoint_node = static_cast(json_parse_size(json, "enterpoint_node")); metadata.base_degree = json_parse_size(json, "base_degree"); metadata.levels_bytes = json_parse_size(json, "levels_bytes"); + metadata.base_nodes_bytes = json_parse_size(json, "base_nodes_bytes"); metadata.base_link_row_bytes = json_parse_size(json, "base_link_row_bytes"); metadata.base_links_bytes = json_parse_size(json, "base_links_bytes"); metadata.upper_nodes_count = json_parse_size(json, "upper_nodes_count"); @@ -766,13 +770,15 @@ inline auto open_layered_dataset_file(const std::string& path) -> npy_file template void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index& index_, const cuvs::util::file_descriptor& output_fd, - size_t output_offset, + size_t base_nodes_offset, + size_t base_links_offset, size_t base_link_row_bytes, size_t maxM0) { - const auto total_start_time = std::chrono::steady_clock::now(); - const auto& graph_fd_opt = index_.graph_fd(); - const auto& mapping_fd_opt = index_.mapping_fd(); + static_assert(std::is_same_v, + "Layered HNSW artifacts store topology ids as uint32_t"); + const auto& graph_fd_opt = index_.graph_fd(); + const auto& mapping_fd_opt = index_.mapping_fd(); RAFT_EXPECTS(graph_fd_opt.has_value() && graph_fd_opt->is_valid(), "Graph file descriptor is not available"); RAFT_EXPECTS(mapping_fd_opt.has_value() && mapping_fd_opt->is_valid(), @@ -812,7 +818,6 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index reordered_to_original(n_rows); - std::vector original_to_reordered(n_rows); cuvs::util::read_large_file( mapping_npy.fd, reordered_to_original.data(), mapping_bytes, mapping_npy.header_size); for (size_t reordered_id = 0; reordered_id < n_rows; ++reordered_id) { @@ -821,7 +826,6 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index(reordered_id); } const auto mapping_elapsed_ms = elapsed_ms_since(mapping_start_time); RAFT_LOG_INFO("HNSW remap: mapping loaded in %ld ms (%.2f GiB, %.2f GiB/s)", @@ -829,112 +833,80 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index(1, target_batch_bytes / base_link_row_bytes); - const size_t max_gap_rows = std::max(1, (128 * 1024) / graph_row_bytes); - const size_t max_span_rows = std::max(1, (8 * 1024 * 1024) / graph_row_bytes); - - RAFT_LOG_INFO( - "HNSW remap: n_rows=%zu degree=%zu maxM0=%zu batch_size=%zu max_gap=%zu rows max_span=%zu rows", - n_rows, - degree, - maxM0, - batch_size, - max_gap_rows, - max_span_rows); - - auto span_buffer = raft::make_host_matrix(static_cast(max_span_rows), - static_cast(degree)); - std::vector output_buffer(batch_size * base_link_row_bytes); - std::vector source_rows(batch_size); - std::vector order(batch_size); - - const auto graph_start_time = std::chrono::steady_clock::now(); - size_t spans_read = 0; - size_t bytes_read = 0; - size_t bytes_written = 0; - - for (size_t batch_start = 0; batch_start < n_rows; batch_start += batch_size) { - const auto current_batch_size = std::min(batch_size, n_rows - batch_start); - std::fill( - output_buffer.begin(), output_buffer.begin() + current_batch_size * base_link_row_bytes, 0); - for (size_t i = 0; i < current_batch_size; ++i) { - source_rows[i] = original_to_reordered[batch_start + i]; - order[i] = i; - } - std::sort(order.begin(), order.begin() + current_batch_size, [&](size_t lhs, size_t rhs) { - return source_rows[lhs] < source_rows[rhs]; - }); - - size_t sorted_pos = 0; - while (sorted_pos < current_batch_size) { - const auto span_start = static_cast(source_rows[order[sorted_pos]]); - auto span_end = span_start; - auto span_next = sorted_pos + 1; - while (span_next < current_batch_size) { - const auto next_row = static_cast(source_rows[order[span_next]]); - if (next_row - span_start + 1 > max_span_rows) { break; } - if (next_row > span_end + max_gap_rows + 1) { break; } - span_end = next_row; - ++span_next; - } - - const auto span_rows = span_end - span_start + 1; - const auto span_bytes = span_rows * graph_row_bytes; - cuvs::util::read_large_file(graph_npy.fd, - span_buffer.data_handle(), - span_bytes, - graph_npy.header_size + span_start * graph_row_bytes); - ++spans_read; - bytes_read += span_bytes; - - for (size_t i = sorted_pos; i < span_next; ++i) { - const auto dst_row = order[i]; - const auto source_row = static_cast(source_rows[dst_row]); - auto* src = span_buffer.data_handle() + (source_row - span_start) * degree; - auto* dst_row_ptr = output_buffer.data() + dst_row * base_link_row_bytes; - hnswlib::linklistsizeint list_count = static_cast(degree); - std::memcpy(dst_row_ptr, &list_count, sizeof(list_count)); - auto* dst = reinterpret_cast(dst_row_ptr + sizeof(hnswlib::linklistsizeint)); - for (size_t j = 0; j < degree; ++j) { - const auto neighbor = static_cast(src[j]); - RAFT_EXPECTS(neighbor < n_rows, - "Invalid reordered neighbor id %zu in ACE graph row %zu", - neighbor, - source_row); - dst[j] = reordered_to_original[neighbor]; + const auto total_start_time = std::chrono::steady_clock::now(); + const auto graph_row_bytes = degree * sizeof(IdxT); + const auto base_node_row_bytes = sizeof(IdxT); + const auto base_topology_row_size = graph_row_bytes + base_node_row_bytes + base_link_row_bytes; + const size_t target_batch_bytes = 64 * 1024 * 1024; + const size_t batch_size = std::max(1, target_batch_bytes / base_topology_row_size); + auto graph_buffer = raft::make_host_matrix(static_cast(batch_size), + static_cast(degree)); + std::vector base_node_buffer(batch_size); + std::vector base_link_buffer(batch_size * base_link_row_bytes); + + size_t graph_bytes_read = 0; + size_t node_bytes_written = 0; + size_t link_bytes_written = 0; + for (size_t source_start = 0; source_start < n_rows; source_start += batch_size) { + const auto current_batch_size = std::min(batch_size, n_rows - source_start); + const auto batch_bytes = current_batch_size * graph_row_bytes; + cuvs::util::read_large_file(graph_npy.fd, + graph_buffer.data_handle(), + batch_bytes, + graph_npy.header_size + source_start * graph_row_bytes); + graph_bytes_read += batch_bytes; + + bool invalid_neighbor = false; +#pragma omp parallel for reduction(|| : invalid_neighbor) + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { + const auto source_row = source_start + static_cast(batch_idx); + const auto original_id = static_cast(reordered_to_original[source_row]); + base_node_buffer[batch_idx] = static_cast(original_id); + auto* dst_row_ptr = base_link_buffer.data() + batch_idx * base_link_row_bytes; + hnswlib::linklistsizeint list_count = static_cast(degree); + std::memcpy(dst_row_ptr, &list_count, sizeof(list_count)); + auto* dst = reinterpret_cast(dst_row_ptr + sizeof(hnswlib::linklistsizeint)); + auto* src = graph_buffer.data_handle() + batch_idx * degree; + for (size_t j = 0; j < degree; ++j) { + const auto neighbor = static_cast(src[j]); + if (neighbor >= n_rows) { + invalid_neighbor = true; + continue; } + dst[j] = reordered_to_original[neighbor]; + } + const auto written_bytes = sizeof(hnswlib::linklistsizeint) + degree * sizeof(IdxT); + if (written_bytes < base_link_row_bytes) { + std::memset(dst_row_ptr + written_bytes, 0, base_link_row_bytes - written_bytes); } - sorted_pos = span_next; } + RAFT_EXPECTS(!invalid_neighbor, "Invalid reordered neighbor id in ACE graph"); - const auto batch_bytes = current_batch_size * base_link_row_bytes; + const auto current_node_bytes = current_batch_size * sizeof(IdxT); + const auto current_link_bytes = current_batch_size * base_link_row_bytes; + cuvs::util::write_large_file(output_fd, + base_node_buffer.data(), + current_node_bytes, + base_nodes_offset + source_start * sizeof(IdxT)); cuvs::util::write_large_file(output_fd, - output_buffer.data(), - batch_bytes, - output_offset + batch_start * base_link_row_bytes); - bytes_written += batch_bytes; + base_link_buffer.data(), + current_link_bytes, + base_links_offset + source_start * base_link_row_bytes); + node_bytes_written += current_node_bytes; + link_bytes_written += current_link_bytes; } - const auto graph_elapsed_ms = elapsed_ms_since(graph_start_time); const auto total_elapsed_ms = elapsed_ms_since(total_start_time); - const auto read_gib = to_gib(bytes_read); - const auto written_gib = to_gib(bytes_written); - const auto ideal_read_bytes = n_rows * graph_row_bytes; - const auto amplification = - ideal_read_bytes == 0 ? 0.0 : static_cast(bytes_read) / ideal_read_bytes; RAFT_LOG_INFO( - "HNSW remap: base links written in %ld ms: read %.2f GiB at %.2f GiB/s across %zu spans " - "(%.2fx amplification), wrote %.2f GiB at %.2f GiB/s", - graph_elapsed_ms, - read_gib, - throughput_gib_per_s(bytes_read, graph_elapsed_ms), - spans_read, - amplification, - written_gib, - throughput_gib_per_s(bytes_written, graph_elapsed_ms)); - RAFT_LOG_INFO("HNSW remap: completed in %ld ms total", total_elapsed_ms); + "HNSW remap: base topology written in %ld ms (graph_read=%.2f GiB %.2f GiB/s, " + "artifact_write=%.2f GiB %.2f GiB/s, nodes=%.2f GiB links=%.2f GiB)", + total_elapsed_ms, + to_gib(graph_bytes_read), + throughput_gib_per_s(graph_bytes_read, total_elapsed_ms), + to_gib(node_bytes_written + link_bytes_written), + throughput_gib_per_s(node_bytes_written + link_bytes_written, total_elapsed_ms), + to_gib(node_bytes_written), + to_gib(link_bytes_written)); } template @@ -975,7 +947,7 @@ auto serialize_to_layered_hnsw_from_disk( RAFT_LOG_INFO("Layered HNSW: generating hierarchy levels"); const auto hierarchy_start_time = std::chrono::steady_clock::now(); - auto hierarchy = make_random_hnsw_level_plan(n_rows, *appr_algo, "Layered HNSW: Level "); + auto hierarchy = make_random_hnsw_level_plan(n_rows, *appr_algo, nullptr); const auto hierarchy_elapsed_ms = elapsed_ms_since(hierarchy_start_time); RAFT_LOG_INFO("Layered HNSW: hierarchy levels generated in %ld ms (max_level=%d, promoted=%zu)", hierarchy_elapsed_ms, @@ -994,6 +966,7 @@ auto serialize_to_layered_hnsw_from_disk( metadata.enterpoint_node = static_cast(hierarchy.order.back()); metadata.base_degree = static_cast(graph_degree_int); metadata.levels_bytes = n_rows * sizeof(uint8_t); + metadata.base_nodes_bytes = n_rows * sizeof(IdxT); metadata.base_link_row_bytes = appr_algo->size_links_level0_; metadata.base_links_bytes = n_rows * metadata.base_link_row_bytes; metadata.upper_link_row_bytes = appr_algo->size_links_per_element_; @@ -1022,7 +995,8 @@ auto serialize_to_layered_hnsw_from_disk( const auto payload_offset = align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); const auto levels_offset = payload_offset; - const auto base_links_offset = levels_offset + metadata.levels_bytes; + const auto base_nodes_offset = levels_offset + metadata.levels_bytes; + const auto base_links_offset = base_nodes_offset + metadata.base_nodes_bytes; const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; const auto final_file_size = upper_links_offset + metadata.upper_links_bytes; @@ -1048,8 +1022,12 @@ auto serialize_to_layered_hnsw_from_disk( RAFT_LOG_INFO("Layered HNSW: writing hnswlib-ready base links"); const auto layer0_start_time = std::chrono::steady_clock::now(); - write_layered_base_links_from_disk( - index_, artifact_fd, base_links_offset, metadata.base_link_row_bytes, metadata.maxM0); + write_layered_base_links_from_disk(index_, + artifact_fd, + base_nodes_offset, + base_links_offset, + metadata.base_link_row_bytes, + metadata.maxM0); const auto layer0_elapsed_ms = elapsed_ms_since(layer0_start_time); RAFT_LOG_INFO("Layered HNSW: base links written in %ld ms (%.2f GiB, %.2f GiB/s effective)", layer0_elapsed_ms, @@ -2081,6 +2059,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, const auto metadata_elapsed_ms = elapsed_ms_since(metadata_start_time); RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); + RAFT_EXPECTS(metadata.dim > 0, "Layered HNSW artifact must contain at least one dimension"); RAFT_EXPECTS(static_cast(dim) == metadata.dim, "Layered HNSW artifact dim (%zu) does not match requested dim (%d)", metadata.dim, @@ -2096,7 +2075,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, const auto payload_offset = align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); const auto levels_offset = payload_offset; - const auto base_links_offset = levels_offset + metadata.levels_bytes; + const auto base_nodes_offset = levels_offset + metadata.levels_bytes; + const auto base_links_offset = base_nodes_offset + metadata.base_nodes_bytes; const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; const auto expected_file_size = upper_links_offset + metadata.upper_links_bytes; @@ -2119,6 +2099,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, "Layered HNSW dataset shape mismatch"); RAFT_EXPECTS(metadata.levels_bytes == metadata.n_rows * sizeof(uint8_t), "Layered HNSW levels section size mismatch"); + RAFT_EXPECTS(metadata.base_nodes_bytes == metadata.n_rows * sizeof(uint32_t), + "Layered HNSW base node section size mismatch"); RAFT_EXPECTS(metadata.base_links_bytes == metadata.n_rows * metadata.base_link_row_bytes, "Layered HNSW base links section size mismatch"); RAFT_EXPECTS(metadata.upper_nodes_bytes == metadata.upper_nodes_count * sizeof(uint32_t), @@ -2172,36 +2154,26 @@ auto deserialize_layered_hnsw(raft::resources const& res, auto num_threads = params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; - RAFT_LOG_INFO("Layered HNSW: initializing in-memory hnswlib index from local dataset"); const size_t target_batch_bytes = 64 * 1024 * 1024; - const size_t row_bytes = metadata.dim * sizeof(T) + metadata.base_link_row_bytes; - const size_t batch_size = std::max(1, target_batch_bytes / row_bytes); + const size_t dataset_row_bytes = metadata.dim * sizeof(T); + const size_t batch_size = std::max(1, target_batch_bytes / dataset_row_bytes); auto dataset_buffer = raft::make_host_matrix(static_cast(batch_size), metadata.dim); - std::vector base_link_buffer(batch_size * metadata.base_link_row_bytes); const auto base_start_time = std::chrono::steady_clock::now(); std::chrono::steady_clock::duration dataset_read_time{}; - std::chrono::steady_clock::duration base_link_read_time{}; std::chrono::steady_clock::duration base_copy_time{}; - size_t dataset_bytes_read = 0; - size_t base_link_bytes_read = 0; + size_t dataset_bytes_read = 0; for (size_t batch_start = 0; batch_start < metadata.n_rows; batch_start += batch_size) { - const auto current_batch_size = std::min(batch_size, metadata.n_rows - batch_start); - auto batch_timer = std::chrono::steady_clock::now(); + const auto current_batch_size = std::min(batch_size, metadata.n_rows - batch_start); + const auto current_dataset_bytes = current_batch_size * metadata.dim * sizeof(T); + auto batch_timer = std::chrono::steady_clock::now(); cuvs::util::read_large_file(dataset_file.fd, dataset_buffer.data_handle(), - current_batch_size * metadata.dim * sizeof(T), + current_dataset_bytes, dataset_file.header_size + batch_start * metadata.dim * sizeof(T)); dataset_read_time += std::chrono::steady_clock::now() - batch_timer; - dataset_bytes_read += current_batch_size * metadata.dim * sizeof(T); - batch_timer = std::chrono::steady_clock::now(); - cuvs::util::read_large_file(artifact_fd, - base_link_buffer.data(), - current_batch_size * metadata.base_link_row_bytes, - base_links_offset + batch_start * metadata.base_link_row_bytes); - base_link_read_time += std::chrono::steady_clock::now() - batch_timer; - base_link_bytes_read += current_batch_size * metadata.base_link_row_bytes; + dataset_bytes_read += current_dataset_bytes; bool link_list_allocation_failed = false; batch_timer = std::chrono::steady_clock::now(); @@ -2210,10 +2182,6 @@ auto deserialize_layered_hnsw(raft::resources const& res, const auto i = batch_start + static_cast(batch_idx); auto level = static_cast(levels_u8[i]); appr_algo->element_levels_[i] = level; - auto ll0 = appr_algo->get_linklist0(i); - memcpy(ll0, - base_link_buffer.data() + batch_idx * metadata.base_link_row_bytes, - metadata.base_link_row_bytes); memcpy(appr_algo->getDataByInternalId(i), dataset_buffer.data_handle() + batch_idx * metadata.dim, appr_algo->data_size_); @@ -2225,7 +2193,6 @@ auto deserialize_layered_hnsw(raft::resources const& res, link_list_allocation_failed = true; continue; } - memset(appr_algo->linkLists_[i], 0, link_list_size); } } if (link_list_allocation_failed) { @@ -2239,36 +2206,83 @@ auto deserialize_layered_hnsw(raft::resources const& res, } const auto base_elapsed_ms = elapsed_ms_since(base_start_time); RAFT_LOG_INFO( - "Layered HNSW load: base layer initialized in %ld ms " - "(dataset read %.2f ms %.2f GiB %.2f GiB/s, links read %.2f ms %.2f GiB %.2f GiB/s, " - "copy %.2f ms)", + "Layered HNSW load: hnswlib data and levels initialized in %ld ms " + "(dataset read %.2f ms %.2f GiB %.2f GiB/s, copy %.2f ms)", base_elapsed_ms, elapsed_ms(dataset_read_time), to_gib(dataset_bytes_read), throughput_gib_per_s(dataset_bytes_read, dataset_read_time), - elapsed_ms(base_link_read_time), - to_gib(base_link_bytes_read), - throughput_gib_per_s(base_link_bytes_read, base_link_read_time), elapsed_ms(base_copy_time)); - RAFT_LOG_INFO("Layered HNSW: loading upper-layer topology"); + const auto base_topology_start_time = std::chrono::steady_clock::now(); + std::chrono::steady_clock::duration base_topology_read_time{}; + std::chrono::steady_clock::duration base_topology_copy_time{}; + size_t base_topology_bytes_read = 0; + const size_t base_topology_row_bytes = sizeof(uint32_t) + metadata.base_link_row_bytes; + const size_t base_topology_batch_size = + std::max(1, target_batch_bytes / base_topology_row_bytes); + std::vector base_node_buffer(base_topology_batch_size); + std::vector base_link_buffer(base_topology_batch_size * metadata.base_link_row_bytes); + for (size_t batch_start = 0; batch_start < metadata.n_rows; + batch_start += base_topology_batch_size) { + const auto current_batch_size = + std::min(base_topology_batch_size, metadata.n_rows - batch_start); + const auto current_base_topology_bytes = + current_batch_size * (sizeof(uint32_t) + metadata.base_link_row_bytes); + auto batch_timer = std::chrono::steady_clock::now(); + cuvs::util::read_large_file(artifact_fd, + base_node_buffer.data(), + current_batch_size * sizeof(uint32_t), + base_nodes_offset + batch_start * sizeof(uint32_t)); + cuvs::util::read_large_file(artifact_fd, + base_link_buffer.data(), + current_batch_size * metadata.base_link_row_bytes, + base_links_offset + batch_start * metadata.base_link_row_bytes); + base_topology_read_time += std::chrono::steady_clock::now() - batch_timer; + base_topology_bytes_read += current_base_topology_bytes; + + batch_timer = std::chrono::steady_clock::now(); + bool invalid_node_id = false; +#pragma omp parallel for num_threads(num_threads) reduction(|| : invalid_node_id) + for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { + const auto node_id = static_cast(base_node_buffer[batch_idx]); + if (node_id >= metadata.n_rows) { + invalid_node_id = true; + continue; + } + auto ll0 = appr_algo->get_linklist0(node_id); + memcpy(ll0, + base_link_buffer.data() + batch_idx * metadata.base_link_row_bytes, + metadata.base_link_row_bytes); + } + RAFT_EXPECTS(!invalid_node_id, "Invalid base-layer node id in layered HNSW artifact"); + base_topology_copy_time += std::chrono::steady_clock::now() - batch_timer; + } + const auto base_topology_elapsed_ms = elapsed_ms_since(base_topology_start_time); + RAFT_LOG_INFO( + "Layered HNSW load: base-layer topology loaded in %ld ms " + "(read %.2f ms %.2f GiB %.2f GiB/s, scatter-copy %.2f ms)", + base_topology_elapsed_ms, + elapsed_ms(base_topology_read_time), + to_gib(base_topology_bytes_read), + throughput_gib_per_s(base_topology_bytes_read, base_topology_read_time), + elapsed_ms(base_topology_copy_time)); + const auto upper_start_time = std::chrono::steady_clock::now(); std::chrono::steady_clock::duration upper_read_time{}; - std::chrono::steady_clock::duration upper_validate_time{}; std::chrono::steady_clock::duration upper_copy_time{}; size_t upper_bytes_read = 0; for (const auto& layer : metadata.layers) { - const auto layer_start_time = std::chrono::steady_clock::now(); - RAFT_LOG_INFO("Layered HNSW: loading layer %zu (%zu rows, degree %zu)", - layer.level, - layer.row_count, - layer.degree); - const auto layer_batch_size = std::min(batch_size, layer.row_count); + const auto layer_row_bytes = sizeof(uint32_t) + metadata.upper_link_row_bytes; + const auto layer_batch_size = + std::min(std::max(1, target_batch_bytes / layer_row_bytes), layer.row_count); std::vector node_buffer(layer_batch_size); std::vector link_buffer(layer_batch_size * metadata.upper_link_row_bytes); for (size_t batch_start = 0; batch_start < layer.row_count; batch_start += layer_batch_size) { const auto current_batch_size = std::min(layer_batch_size, layer.row_count - batch_start); - auto batch_timer = std::chrono::steady_clock::now(); + const auto current_upper_bytes = + current_batch_size * (sizeof(uint32_t) + metadata.upper_link_row_bytes); + auto batch_timer = std::chrono::steady_clock::now(); cuvs::util::read_large_file( artifact_fd, node_buffer.data(), @@ -2280,45 +2294,42 @@ auto deserialize_layered_hnsw(raft::resources const& res, current_batch_size * metadata.upper_link_row_bytes, upper_links_offset + (layer.link_offset + batch_start) * metadata.upper_link_row_bytes); upper_read_time += std::chrono::steady_clock::now() - batch_timer; - upper_bytes_read += current_batch_size * (sizeof(uint32_t) + metadata.upper_link_row_bytes); - batch_timer = std::chrono::steady_clock::now(); + upper_bytes_read += current_upper_bytes; + batch_timer = std::chrono::steady_clock::now(); + bool invalid_node_id = false; + bool invalid_node_level = false; +#pragma omp parallel for num_threads(num_threads) \ + reduction(|| : invalid_node_id, invalid_node_level) for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { const auto node_id = static_cast(node_buffer[batch_idx]); - RAFT_EXPECTS(node_id < metadata.n_rows, - "Invalid upper-layer node id %zu in layered HNSW artifact", - node_id); - RAFT_EXPECTS(layer.level <= static_cast(levels_u8[node_id]), - "Layered HNSW artifact references node %zu at invalid level %zu", - node_id, - layer.level); - } - upper_validate_time += std::chrono::steady_clock::now() - batch_timer; - batch_timer = std::chrono::steady_clock::now(); -#pragma omp parallel for num_threads(num_threads) - for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); - ++batch_idx) { - const auto node_id = static_cast(node_buffer[batch_idx]); - auto ll = appr_algo->get_linklist(node_id, layer.level); + if (node_id >= metadata.n_rows) { + invalid_node_id = true; + continue; + } + if (layer.level > static_cast(levels_u8[node_id])) { + invalid_node_level = true; + continue; + } + auto ll = appr_algo->get_linklist(node_id, layer.level); memcpy(ll, link_buffer.data() + batch_idx * metadata.upper_link_row_bytes, metadata.upper_link_row_bytes); } + RAFT_EXPECTS(!invalid_node_id, "Invalid upper-layer node id in layered HNSW artifact"); + RAFT_EXPECTS(!invalid_node_level, + "Layered HNSW artifact references a node at an invalid upper level"); upper_copy_time += std::chrono::steady_clock::now() - batch_timer; } - RAFT_LOG_INFO("Layered HNSW load: layer %zu loaded in %ld ms", - layer.level, - elapsed_ms_since(layer_start_time)); } const auto upper_elapsed_ms = elapsed_ms_since(upper_start_time); RAFT_LOG_INFO( "Layered HNSW load: upper layers loaded in %ld ms " - "(read %.2f ms %.2f GiB %.2f GiB/s, validate %.2f ms, copy %.2f ms)", + "(read %.2f ms %.2f GiB %.2f GiB/s, validate+copy %.2f ms)", upper_elapsed_ms, elapsed_ms(upper_read_time), to_gib(upper_bytes_read), throughput_gib_per_s(upper_bytes_read, upper_read_time), - elapsed_ms(upper_validate_time), elapsed_ms(upper_copy_time)); hnsw_index->set_index(std::move(appr_algo)); diff --git a/examples/cpp/src/hnsw_ace_layered_example.cu b/examples/cpp/src/hnsw_ace_layered_example.cu index 02f735ef37..f11551df20 100644 --- a/examples/cpp/src/hnsw_ace_layered_example.cu +++ b/examples/cpp/src/hnsw_ace_layered_example.cu @@ -23,16 +23,16 @@ // index_dir/hnsw_index.cuvs // fixed header + metadata JSON // levels: uint8 [N], max HNSW level for each original row id -// base links: hnswlib-ready [count + uint32 neighbors padded to maxM0] +// base nodes + base links: uint32 node ids with hnswlib-ready link rows // upper nodes + upper links: hnswlib-ready upper-layer topology // // The transferred index artifact is topology-only. The dataset is loaded locally during // deserialization from hnsw::index_params::dataset_path. The loader supports .npy and ANN benchmark // *.bin datasets; this example writes a local dataset .npy only to make the demo self-contained. // -// Layer 0 rows and neighbor IDs are original dataset row IDs. Upper layers are generated with the -// same level/order/KNN logic as serialize_to_hnswlib_from_disk, then stored as hnswlib-ready link -// rows so deserialization does no graph remapping or link padding on the search node. +// Layer 0 node IDs and neighbor IDs are original dataset row IDs. Upper layers are generated with +// the same level/order/KNN logic as serialize_to_hnswlib_from_disk, then stored as hnswlib-ready +// link rows so deserialization does no graph remapping or link padding on the search node. #include #include From 4514bc85b1c1743ee05af9eb4b2465967925c390 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Mon, 1 Jun 2026 10:15:11 +0200 Subject: [PATCH 06/12] Unify the ACE logging format --- .../neighbors/detail/cagra/cagra_build.cuh | 234 +++++++++--------- cpp/src/neighbors/detail/hnsw.hpp | 160 ++++++++---- 2 files changed, 238 insertions(+), 156 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 96ff8344d3..3edbec46af 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -51,6 +51,13 @@ namespace cuvs::neighbors::cagra::detail { constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } +inline auto progress_step_10(size_t completed, size_t total) -> size_t +{ + if (total == 0) { return 100; } + const auto percent = (std::min(completed, total) * 100) / total; + return completed >= total ? 100 : (percent / 10) * 10; +} + template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { @@ -99,7 +106,7 @@ void ace_get_partition_labels( const size_t min_samples = 100 * n_partitions; n_samples = std::max(n_samples, min_samples); n_samples = std::min(n_samples, dataset_size); - RAFT_LOG_DEBUG("ACE: n_samples: %lu", n_samples); + RAFT_LOG_DEBUG("ACE build: partition labeling uses %lu sampled vectors", n_samples); auto sample_db = raft::make_host_matrix(n_samples, dataset_dim); #pragma omp parallel for @@ -123,18 +130,10 @@ void ace_get_partition_labels( auto _sub_distances = raft::make_host_matrix(chunk_size, n_partitions); auto _sub_dataset_dev = raft::make_device_matrix(res, chunk_size, dataset_dim); auto _sub_distances_dev = raft::make_device_matrix(res, chunk_size, n_partitions); - size_t report_interval = dataset_size / 10; - report_interval = (report_interval / chunk_size) * chunk_size; - report_interval = std::max(report_interval, chunk_size); + size_t next_progress_percent = 10; for (size_t i_base = 0; i_base < dataset_size; i_base += chunk_size) { const size_t sub_dataset_size = std::min(chunk_size, dataset_size - i_base); - if (i_base % report_interval == 0) { - RAFT_LOG_INFO("ACE: Processing chunk %lu / %lu (%.1f%%)", - i_base, - dataset_size, - static_cast(100 * i_base) / dataset_size); - } auto sub_dataset = raft::make_host_matrix_view( _sub_dataset.data_handle(), sub_dataset_size, dataset_dim); @@ -191,6 +190,16 @@ void ace_get_partition_labels( #pragma omp atomic update partition_histogram(augmented_label, 1) += 1; } + + const auto completed_rows = i_base + sub_dataset_size; + const auto report_percent = progress_step_10(completed_rows, dataset_size); + if (report_percent >= next_progress_percent) { + RAFT_LOG_INFO("ACE build: partition labeling progress %zu%% (%zu/%zu rows)", + report_percent, + completed_rows, + dataset_size); + next_progress_percent = report_percent + 10; + } } } @@ -261,15 +270,15 @@ void ace_check_partition_sizes( if (total_count > 0 && total_count < very_small_threshold) { RAFT_LOG_WARN( - "ACE: Partition %lu is very small (%lu vectors, expected ~%.1f). This may affect graph " - "quality.", + "ACE build: partition %lu is very small (%lu vectors, expected ~%.1f); graph quality may " + "be affected", c, total_count, expected_avg_vectors); } else if (total_count > very_large_threshold) { RAFT_LOG_WARN( - "ACE: Partition %lu is very large (%lu vectors, expected ~%.1f, threshold: %lu). This may " - "indicate imbalance and can lead to memory issues in restricted environments.", + "ACE build: partition %lu is very large (%lu vectors, expected ~%.1f, threshold=%lu); " + "partition imbalance may increase memory pressure", c, total_count, expected_avg_vectors, @@ -482,11 +491,10 @@ void ace_reorder_and_store_dataset( size_t dataset_dim = dataset.extent(1); size_t n_partitions = partition_histogram.extent(0); - RAFT_LOG_DEBUG( - "ACE: Reordering and storing dataset to disk (%lu vectors, %lu dimensions, %lu partitions)", - dataset_size, - dataset_dim, - n_partitions); + RAFT_LOG_DEBUG("ACE build: reordering dataset to disk (rows=%lu dim=%lu partitions=%lu)", + dataset_size, + dataset_dim, + n_partitions); // Calculate total sizes for pre-allocation size_t total_core_vectors = 0; @@ -509,10 +517,10 @@ void ace_reorder_and_store_dataset( size_t reordered_file_size = total_core_vectors * vector_size; size_t augmented_file_size = total_augmented_vectors * vector_size; - RAFT_LOG_DEBUG("ACE: Reordered dataset: %lu core vectors (%.2f GiB)", + RAFT_LOG_DEBUG("ACE build: reordered dataset section rows=%lu size=%.2f GiB", total_core_vectors, reordered_file_size / (1024.0 * 1024.0 * 1024.0)); - RAFT_LOG_DEBUG("ACE: Augmented dataset: %lu secondary vectors (%.2f GiB)", + RAFT_LOG_DEBUG("ACE build: augmented dataset section rows=%lu size=%.2f GiB", total_augmented_vectors, augmented_file_size / (1024.0 * 1024.0 * 1024.0)); @@ -541,7 +549,7 @@ void ace_reorder_and_store_dataset( disk_write_size = std::min(disk_write_size, 64 * 1024 * 1024); size_t vectors_per_buffer = std::max(64, disk_write_size / vector_size); - RAFT_LOG_DEBUG("ACE: Reorder buffers: %lu vectors per buffer (%.2f MiB)", + RAFT_LOG_DEBUG("ACE build: reorder buffer rows=%lu size=%.2f MiB", vectors_per_buffer, to_mib(vectors_per_buffer * vector_size)); @@ -593,8 +601,8 @@ void ace_reorder_and_store_dataset( } }; - size_t vectors_processed = 0; - const size_t log_interval = std::max(dataset_size / 10, size_t(1)); + size_t vectors_processed = 0; + size_t next_progress_percent = 10; for (size_t i = 0; i < dataset_size; i++) { size_t core_partition = partition_labels(i, 0); size_t secondary_partition = partition_labels(i, 1); @@ -623,16 +631,18 @@ void ace_reorder_and_store_dataset( } vectors_processed++; - if (vectors_processed % log_interval == 0) { - RAFT_LOG_INFO("ACE: Processed %lu/%lu vectors (%.1f%%)", + const auto report_percent = progress_step_10(vectors_processed, dataset_size); + if (report_percent >= next_progress_percent) { + RAFT_LOG_INFO("ACE build: dataset reorder progress %zu%% (%zu/%zu rows)", + report_percent, vectors_processed, - dataset_size, - 100.0 * vectors_processed / dataset_size); + dataset_size); + next_progress_percent = report_percent + 10; } } // Flush all remaining buffers - RAFT_LOG_DEBUG("ACE: Flushing remaining buffers..."); + RAFT_LOG_DEBUG("ACE build: flushing remaining reorder buffers"); #pragma omp parallel sections { #pragma omp section @@ -662,12 +672,12 @@ void ace_reorder_and_store_dataset( elapsed_ms > 0 ? to_mib(total_bytes_written) / (elapsed_ms / 1000.0) : 0.0; RAFT_LOG_INFO( - "ACE: Dataset (%.2f GiB reordered, %.2f GiB augmented, %.2f GiB mapping) reordering completed " - "in %ld ms (%.1f MiB/s)", + "ACE build: dataset reorder completed in %ld ms (reordered=%.2f GiB augmented=%.2f GiB " + "mapping=%.2f GiB, %.1f MiB/s)", + elapsed_ms, reordered_file_size / (1024.0 * 1024.0 * 1024.0), augmented_file_size / (1024.0 * 1024.0 * 1024.0), mapping_file_size / (1024.0 * 1024.0 * 1024.0), - elapsed_ms, throughput_mb_s); } @@ -685,17 +695,17 @@ void ace_load_partition_dataset_from_disk( { size_t n_partitions = partition_histogram.extent(0); - RAFT_LOG_DEBUG("ACE: Loading partition %lu dataset from disk", partition_id); + RAFT_LOG_DEBUG("ACE build: loading partition %lu dataset from disk", partition_id); size_t core_size = partition_histogram(partition_id, 0); size_t augmented_size = partition_histogram(partition_id, 1); size_t total_partition_size = core_size + augmented_size; - RAFT_LOG_DEBUG("ACE: Partition %lu: %lu core + %lu augmented = %lu total vectors", + RAFT_LOG_DEBUG("ACE build: partition %lu rows=%lu (core=%lu augmented=%lu)", partition_id, + total_partition_size, core_size, - augmented_size, - total_partition_size); + augmented_size); RAFT_EXPECTS(static_cast(sub_dataset.extent(0)) == total_partition_size, "sub_dataset rows (%lu) must match total partition size (%lu)", @@ -712,10 +722,10 @@ void ace_load_partition_dataset_from_disk( const std::string augmented_dataset_path = build_dir + "/augmented_dataset.npy"; if (!std::filesystem::exists(reordered_dataset_path)) { - RAFT_FAIL("ACE: Required file does not exist: %s", reordered_dataset_path.c_str()); + RAFT_FAIL("ACE build: required file does not exist: %s", reordered_dataset_path.c_str()); } if (!std::filesystem::exists(augmented_dataset_path)) { - RAFT_FAIL("ACE: Required file does not exist: %s", augmented_dataset_path.c_str()); + RAFT_FAIL("ACE build: required file does not exist: %s", augmented_dataset_path.c_str()); } size_t core_header_size = 0; @@ -748,7 +758,7 @@ void ace_load_partition_dataset_from_disk( core_file_offset += core_header_size; augmented_file_offset += augmented_header_size; - RAFT_LOG_DEBUG("ACE: Core file offset: %lu bytes, Augmented file offset: %lu bytes", + RAFT_LOG_DEBUG("ACE build: core file offset=%lu bytes augmented file offset=%lu bytes", core_file_offset, augmented_file_offset); @@ -763,7 +773,7 @@ void ace_load_partition_dataset_from_disk( try { if (core_size > 0) { RAFT_LOG_DEBUG( - "ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); + "ACE build: reading %lu core vectors from offset %lu", core_size, core_file_offset); cuvs::util::file_descriptor reordered_fd(reordered_dataset_path, O_RDONLY); const size_t core_bytes = core_size * vector_size; cuvs::util::read_large_file( @@ -777,7 +787,7 @@ void ace_load_partition_dataset_from_disk( { try { if (augmented_size > 0) { - RAFT_LOG_DEBUG("ACE: Reading %lu augmented vectors from offset %lu", + RAFT_LOG_DEBUG("ACE build: reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); cuvs::util::file_descriptor augmented_fd(augmented_dataset_path, O_RDONLY); @@ -879,7 +889,8 @@ bool ace_check_use_disk_mode(bool use_disk, // Use overridden memory limits if provided (> 0), otherwise query actual system memory if (max_host_memory_gb.has_value() && max_host_memory_gb.value() > 0) { mem.available_host_memory = static_cast(max_host_memory_gb.value() * (1ULL << 30)); - RAFT_LOG_INFO("ACE: Using overridden host memory limit: %.2f GiB", max_host_memory_gb.value()); + RAFT_LOG_INFO("ACE build: using overridden host memory limit %.2f GiB", + max_host_memory_gb.value()); } else { mem.available_host_memory = cuvs::util::get_free_host_memory(); } @@ -903,7 +914,7 @@ bool ace_check_use_disk_mode(bool use_disk, mem.total_size = mem.partition_labels_size + mem.id_mapping_size + mem.sub_dataset_size + mem.sub_graph_size + mem.cagra_graph_size; - RAFT_LOG_INFO("ACE: Estimated host memory required: %.2f GiB, available: %.2f GiB", + RAFT_LOG_INFO("ACE build: host memory estimate required=%.2f GiB available=%.2f GiB", to_gib(mem.total_size), to_gib(mem.available_host_memory)); @@ -915,7 +926,8 @@ bool ace_check_use_disk_mode(bool use_disk, // TODO: Extend model or use managed memory if running out of GPU memory. if (max_gpu_memory_gb.has_value() && max_gpu_memory_gb.value() > 0) { mem.available_gpu_memory = static_cast(max_gpu_memory_gb.value() * (1ULL << 30)); - RAFT_LOG_INFO("ACE: Using overridden GPU memory limit: %.2f GiB", max_gpu_memory_gb.value()); + RAFT_LOG_INFO("ACE build: using overridden GPU memory limit %.2f GiB", + max_gpu_memory_gb.value()); } else { mem.available_gpu_memory = rmm::available_device_memory().second; } @@ -923,7 +935,7 @@ bool ace_check_use_disk_mode(bool use_disk, static_cast(usable_gpu_memory_fraction * mem.available_gpu_memory) < std::max(mem.sub_graph_size, mem.sub_dataset_size); - RAFT_LOG_INFO("ACE: Estimated GPU memory required: %.2f GiB, available: %.2f GiB", + RAFT_LOG_INFO("ACE build: GPU memory estimate required=%.2f GiB available=%.2f GiB", to_gib(mem.cagra_graph_size), to_gib(mem.available_gpu_memory)); @@ -934,7 +946,7 @@ bool ace_check_use_disk_mode(bool use_disk, valid_build_dir &= build_dir.find('\0') == std::string::npos; valid_build_dir &= build_dir.find("//") == std::string::npos; if (!valid_build_dir) { - RAFT_LOG_WARN("ACE: Invalid build_dir path, resetting to default: /tmp/ace_build"); + RAFT_LOG_WARN("ACE build: invalid build_dir path, resetting to /tmp/ace_build"); build_dir = "/tmp/ace_build"; } if (mkdir(build_dir.c_str(), 0755) != 0 && errno != EEXIST) { @@ -943,24 +955,19 @@ bool ace_check_use_disk_mode(bool use_disk, } if (host_memory_limited && gpu_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in host and GPU memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); + RAFT_LOG_INFO("ACE build: graph does not fit in host or GPU memory; using disk mode at %s", + build_dir.c_str()); } else if (host_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in host memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); + RAFT_LOG_INFO("ACE build: graph does not fit in host memory; using disk mode at %s", + build_dir.c_str()); } else if (gpu_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in GPU memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); + RAFT_LOG_INFO("ACE build: graph does not fit in GPU memory; using disk mode at %s", + build_dir.c_str()); } else if (use_disk) { - RAFT_LOG_INFO( - "ACE: Graph fits in host and GPU memory but disk mode is forced. Using disk-mode with " - "temporary storage %s", - build_dir.c_str()); + RAFT_LOG_INFO("ACE build: graph fits in host and GPU memory but disk mode is forced; using %s", + build_dir.c_str()); } else { - RAFT_LOG_INFO("ACE: Graph fits in host and GPU memory. Using in-memory mode."); + RAFT_LOG_INFO("ACE build: graph fits in host and GPU memory; using in-memory mode"); } return use_disk_mode; @@ -1005,8 +1012,8 @@ void ace_validate_disk_mode_partitions(size_t& n_partitions, disk_mode_host_required) { host_memory_insufficient = true; RAFT_LOG_WARN( - "ACE: Host memory insufficient for disk mode. Required: %.2f GiB, available: %.2f GiB. " - "Per-partition breakdown: dataset %.2f GiB, graph %.2f GiB, workspace %.2f GiB", + "ACE build: host memory is insufficient for disk mode: required=%.2f GiB available=%.2f " + "GiB (partition dataset=%.2f GiB graph=%.2f GiB workspace=%.2f GiB)", to_gib(disk_mode_host_required), to_gib(mem.available_host_memory), to_gib(mem.sub_dataset_size), @@ -1017,8 +1024,8 @@ void ace_validate_disk_mode_partitions(size_t& n_partitions, double available_for_scaling = usable_cpu_memory_fraction * mem.available_host_memory - mem.partition_labels_size - mem.id_mapping_size; RAFT_EXPECTS(available_for_scaling > 0, - "ACE: Host memory insufficient even for constant overhead (labels + id_mapping). " - "Required: %.2f GiB, available: %.2f GiB", + "ACE build: host memory is insufficient even for constant overhead " + "(labels + id_mapping): required=%.2f GiB available=%.2f GiB", to_gib(mem.partition_labels_size + mem.id_mapping_size), to_gib(usable_cpu_memory_fraction * mem.available_host_memory)); host_suggested_partitions = static_cast( @@ -1040,9 +1047,8 @@ void ace_validate_disk_mode_partitions(size_t& n_partitions, disk_mode_gpu_required) { gpu_memory_insufficient = true; RAFT_LOG_WARN( - "ACE: GPU memory insufficient for per-partition processing. Required: %.2f GiB, " - "available: %.2f GiB. Per-partition breakdown: dataset %.2f GiB, graph %.2f GiB, " - "workspace %.2f GiB", + "ACE build: GPU memory is insufficient for partition processing: required=%.2f GiB " + "available=%.2f GiB (partition dataset=%.2f GiB graph=%.2f GiB workspace=%.2f GiB)", to_gib(disk_mode_gpu_required), to_gib(mem.available_gpu_memory), to_gib(mem.sub_dataset_size), @@ -1060,8 +1066,7 @@ void ace_validate_disk_mode_partitions(size_t& n_partitions, size_t new_n_partitions = std::max(host_suggested_partitions, gpu_suggested_partitions); RAFT_LOG_WARN( - "ACE: Automatically increasing number of partitions from %zu to %zu to satisfy memory " - "constraints.%s%s", + "ACE build: increasing partitions from %zu to %zu to satisfy memory constraints.%s%s", original_n_partitions, new_n_partitions, host_memory_insufficient @@ -1088,8 +1093,8 @@ void ace_validate_disk_mode_partitions(size_t& n_partitions, guarantee_connectivity); RAFT_LOG_INFO( - "ACE: Updated per-partition memory estimates: dataset %.2f GiB, graph %.2f GiB, " - "host workspace %.2f GiB, GPU workspace %.2f GiB", + "ACE build: updated partition memory estimates dataset=%.2f GiB graph=%.2f GiB host " + "workspace=%.2f GiB GPU workspace=%.2f GiB", to_gib(mem.sub_dataset_size), to_gib(mem.sub_graph_size), to_gib(new_opt_host_ws), @@ -1131,15 +1136,15 @@ index build_ace(raft::resources const& res, size_t dataset_size = dataset.extent(0); size_t dataset_dim = dataset.extent(1); - RAFT_EXPECTS(dataset_size > 0, "ACE: Dataset must not be empty"); + RAFT_EXPECTS(dataset_size > 0, "ACE build: dataset must not be empty"); if (dataset_size < 1000) { - RAFT_LOG_WARN("ACE: Very small dataset size (%zu), consider using regular CAGRA build instead.", + RAFT_LOG_WARN("ACE build: very small dataset size (%zu); regular CAGRA may be simpler", dataset_size); } - RAFT_EXPECTS(dataset_dim > 0, "ACE: Dataset dimension must be greater than 0"); + RAFT_EXPECTS(dataset_dim > 0, "ACE build: dataset dimension must be greater than 0"); RAFT_EXPECTS(params.intermediate_graph_degree > 0, - "ACE: Intermediate graph degree must be greater than 0"); - RAFT_EXPECTS(params.graph_degree > 0, "ACE: Graph degree must be greater than 0"); + "ACE build: intermediate graph degree must be greater than 0"); + RAFT_EXPECTS(params.graph_degree > 0, "ACE build: graph degree must be greater than 0"); size_t n_partitions = npartitions; if (n_partitions == 0) { @@ -1152,17 +1157,22 @@ index build_ace(raft::resources const& res, n_partitions = dataset_size / min_required_per_partition; if (n_partitions < 2) { RAFT_LOG_WARN( - "ACE: Reduced number of partitions to the minimum of 2 to avoid tiny partitions. Consider " - "using regular CAGRA build instead."); + "ACE build: reduced partitions to the minimum of 2 to avoid tiny partitions; regular " + "CAGRA may be simpler"); n_partitions = 2; } else { - RAFT_LOG_WARN("ACE: Reduced number of partitions to %zu to avoid tiny partitions", - n_partitions); + RAFT_LOG_WARN("ACE build: reduced partitions to %zu to avoid tiny partitions", n_partitions); } } auto total_start = std::chrono::high_resolution_clock::now(); - RAFT_LOG_INFO("ACE: Starting partitioned CAGRA build with %zu partitions", n_partitions); + RAFT_LOG_INFO( + "ACE build: start rows=%zu dim=%zu partitions=%zu graph_degree=%zu intermediate_degree=%zu", + dataset_size, + dataset_dim, + n_partitions, + static_cast(params.graph_degree), + static_cast(params.intermediate_graph_degree)); size_t intermediate_degree = params.intermediate_graph_degree; size_t graph_degree = params.graph_degree; @@ -1228,7 +1238,7 @@ index build_ace(raft::resources const& res, build_dir + "/cagra_graph.npy", {dataset_size, graph_degree}); RAFT_LOG_DEBUG( - "ACE: Wrote numpy headers (reordered: %zu, augmented: %zu, mapping: %zu, graph: %zu bytes)", + "ACE build: wrote numpy headers (reordered=%zu augmented=%zu mapping=%zu graph=%zu bytes)", reordered_header_size, augmented_header_size, mapping_header_size, @@ -1259,11 +1269,9 @@ index build_ace(raft::resources const& res, auto partition_elapsed = std::chrono::duration_cast(partition_end - partition_start) .count(); - RAFT_LOG_INFO( - "ACE: Partition labeling completed in %ld ms (min_partition_size: " - "%lu)", - partition_elapsed, - min_partition_size); + RAFT_LOG_INFO("ACE build: partition labeling completed in %ld ms (min_partition_size=%lu)", + partition_elapsed, + min_partition_size); // Create vector lists for each partition auto vectorlist_start = std::chrono::high_resolution_clock::now(); @@ -1288,7 +1296,7 @@ index build_ace(raft::resources const& res, auto vectorlist_elapsed = std::chrono::duration_cast(vectorlist_end - vectorlist_start) .count(); - RAFT_LOG_INFO("ACE: Vector list creation completed in %ld ms", vectorlist_elapsed); + RAFT_LOG_INFO("ACE build: partition mapping completed in %ld ms", vectorlist_elapsed); // Reorder the dataset based on partitions and store to disk. Uses write buffers to improve // performance. @@ -1319,7 +1327,7 @@ index build_ace(raft::resources const& res, // Process each partition auto partition_processing_start = std::chrono::high_resolution_clock::now(); for (size_t partition_id = 0; partition_id < n_partitions; partition_id++) { - RAFT_LOG_DEBUG("ACE: Processing partition %lu/%lu", partition_id + 1, n_partitions); + RAFT_LOG_DEBUG("ACE build: processing partition %lu/%lu", partition_id + 1, n_partitions); auto start = std::chrono::high_resolution_clock::now(); // Extract vectors for this partition @@ -1328,10 +1336,11 @@ index build_ace(raft::resources const& res, size_t sub_dataset_size = core_sub_dataset_size + augmented_sub_dataset_size; if (sub_dataset_size == 0) { - RAFT_LOG_WARN("ACE: Skipping empty partition %lu", partition_id); + RAFT_LOG_WARN("ACE build: Skipping empty partition %lu", partition_id); continue; } - RAFT_LOG_DEBUG("ACE: Sub-dataset size: %lu (%lu + %lu)", + RAFT_LOG_DEBUG("ACE build: partition %lu rows=%lu (core=%lu augmented=%lu)", + partition_id, sub_dataset_size, core_sub_dataset_size, augmented_sub_dataset_size); @@ -1444,12 +1453,12 @@ index build_ace(raft::resources const& res, ? to_mib(core_sub_dataset_size * dataset_dim * sizeof(T)) / (write_elapsed / 1000.0) : 0.0; RAFT_LOG_INFO( - "ACE: Partition %4lu (%8lu + %8lu) completed in %6ld ms: read %6ld ms (%7.1f MiB/s), " - "optimize %6ld ms, adjust %6ld ms, write %6ld ms (%7.1f MiB/s)", + "ACE build: partition %4lu completed in %6ld ms (core=%8lu augmented=%8lu read=%6ld ms " + "%.1f MiB/s optimize=%6ld ms adjust=%6ld ms write=%6ld ms %.1f MiB/s)", partition_id, + elapsed_ms, core_sub_dataset_size, augmented_sub_dataset_size, - elapsed_ms, read_elapsed, read_throughput, optimize_elapsed, @@ -1462,7 +1471,7 @@ index build_ace(raft::resources const& res, auto partition_processing_elapsed = std::chrono::duration_cast( partition_processing_end - partition_processing_start) .count(); - RAFT_LOG_INFO("ACE: All partition processing completed in %ld ms (%zu partitions)", + RAFT_LOG_INFO("ACE build: partition graph build completed in %ld ms (partitions=%zu)", partition_processing_elapsed, n_partitions); @@ -1472,7 +1481,7 @@ index build_ace(raft::resources const& res, const std::string augmented_dataset_path = build_dir + "/augmented_dataset.npy"; if (std::filesystem::exists(augmented_dataset_path)) { std::filesystem::remove(augmented_dataset_path); - RAFT_LOG_INFO("ACE: Removed augmented dataset file to save disk space"); + RAFT_LOG_INFO("ACE build: removed temporary augmented dataset"); } } @@ -1488,12 +1497,12 @@ index build_ace(raft::resources const& res, idx.update_dataset(res, dataset); } catch (std::bad_alloc& e) { RAFT_LOG_WARN( - "Insufficient GPU memory to attach dataset to ACE index. Only the graph will be " - "stored."); + "ACE build: insufficient GPU memory to attach dataset to index; only the graph will " + "be stored"); } catch (raft::logic_error& e) { RAFT_LOG_WARN( - "Insufficient GPU memory to attach dataset to ACE index. Only the graph will be " - "stored."); + "ACE build: insufficient GPU memory to attach dataset to index; only the graph will " + "be stored"); } } } else { @@ -1501,37 +1510,36 @@ index build_ace(raft::resources const& res, idx.update_graph(res, std::move(graph_fd)); idx.update_mapping(res, std::move(mapping_fd)); - RAFT_LOG_INFO( - "ACE: Set disk storage at %s (dataset shape [%zu, %zu], graph shape [%zu, %zu])", - build_dir.c_str(), - idx.size(), - idx.dim(), - idx.size(), - idx.graph_degree()); + RAFT_LOG_INFO("ACE build: disk artifacts ready at %s (dataset=[%zu,%zu] graph=[%zu,%zu])", + build_dir.c_str(), + idx.size(), + idx.dim(), + idx.size(), + idx.graph_degree()); } auto index_creation_end = std::chrono::high_resolution_clock::now(); auto index_creation_elapsed = std::chrono::duration_cast( index_creation_end - index_creation_start) .count(); - RAFT_LOG_INFO("ACE: Final index creation completed in %ld ms", index_creation_elapsed); + RAFT_LOG_INFO("ACE build: final index initialized in %ld ms", index_creation_elapsed); auto total_end = std::chrono::high_resolution_clock::now(); auto total_elapsed = std::chrono::duration_cast(total_end - total_start).count(); - RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed); + RAFT_LOG_INFO("ACE build: completed in %ld ms", total_elapsed); return idx; } catch (const std::exception& e) { // Clean up build directory on failure if we created it - RAFT_LOG_ERROR("ACE: Build failed with exception: %s", e.what()); + RAFT_LOG_ERROR("ACE build: failed with exception: %s", e.what()); if (cleanup_on_failure && !build_dir.empty()) { - RAFT_LOG_INFO("ACE: Cleaning up build directory: %s", build_dir.c_str()); + RAFT_LOG_INFO("ACE build: cleaning up build directory %s", build_dir.c_str()); try { std::filesystem::remove_all(build_dir); - RAFT_LOG_INFO("ACE: Successfully removed build directory"); + RAFT_LOG_INFO("ACE build: removed build directory"); } catch (const std::exception& cleanup_error) { - RAFT_LOG_WARN("ACE: Failed to clean up build directory: %s", cleanup_error.what()); + RAFT_LOG_WARN("ACE build: failed to clean up build directory: %s", cleanup_error.what()); } } // Re-throw the original exception diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index d090de07f8..5be875d804 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -411,6 +411,13 @@ inline auto throughput_gib_per_s(size_t bytes, std::chrono::steady_clock::durati return elapsed_s > 0.0 ? to_gib(bytes) / elapsed_s : 0.0; } +inline auto progress_step_10(size_t completed, size_t total) -> size_t +{ + if (total == 0) { return 100; } + const auto percent = (std::min(completed, total) * 100) / total; + return completed >= total ? 100 : (percent / 10) * 10; +} + inline auto make_hnsw_level_plan_from_levels(size_t n_rows, std::vector&& levels, bool build_reverse_order, @@ -486,7 +493,7 @@ void build_hnsw_upper_layer_graphs( const auto row_count = plan.n_rows - start_idx; const auto neighbor_size = hnsw_upper_layer_degree(row_count, M); - RAFT_LOG_INFO("Compute hierarchy neighbors level %zu", pt_level); + RAFT_LOG_INFO("Layered HNSW artifact: computing upper-layer neighbors level=%zu", pt_level); auto host_neighbors = raft::make_host_matrix( static_cast(row_count), static_cast(neighbor_size)); @@ -814,7 +821,7 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index= sizeof(hnswlib::linklistsizeint) + maxM0 * sizeof(IdxT), "Base link row size is too small"); - RAFT_LOG_INFO("HNSW remap: loading ACE reordered-to-original mapping (%zu rows)", n_rows); + RAFT_LOG_INFO("Layered HNSW artifact: loading ACE row mapping (%zu rows)", n_rows); const auto mapping_start_time = std::chrono::steady_clock::now(); const auto mapping_bytes = n_rows * sizeof(IdxT); std::vector reordered_to_original(n_rows); @@ -828,11 +835,18 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index= next_report_percent) { + const auto elapsed = elapsed_ms_since(total_start_time); + RAFT_LOG_INFO( + "Layered HNSW artifact: base topology progress %zu%% (%zu/%zu rows, read=%.2f GiB " + "write=%.2f GiB, %ld ms)", + report_percent, + completed_rows, + n_rows, + to_gib(graph_bytes_read), + to_gib(node_bytes_written + link_bytes_written), + elapsed); + next_report_percent = report_percent + 10; + } } const auto total_elapsed_ms = elapsed_ms_since(total_start_time); RAFT_LOG_INFO( - "HNSW remap: base topology written in %ld ms (graph_read=%.2f GiB %.2f GiB/s, " + "Layered HNSW artifact: base topology written in %ld ms (graph_read=%.2f GiB %.2f GiB/s, " "artifact_write=%.2f GiB %.2f GiB/s, nodes=%.2f GiB links=%.2f GiB)", total_elapsed_ms, to_gib(graph_bytes_read), @@ -945,14 +983,15 @@ auto serialize_to_layered_hnsw_from_disk( auto appr_algo = std::make_unique::type>>( hnsw_index->get_space(), 1, (graph_degree_int + 1) / 2, params.ef_construction); - RAFT_LOG_INFO("Layered HNSW: generating hierarchy levels"); + RAFT_LOG_INFO("Layered HNSW artifact: generating hierarchy levels"); const auto hierarchy_start_time = std::chrono::steady_clock::now(); auto hierarchy = make_random_hnsw_level_plan(n_rows, *appr_algo, nullptr); const auto hierarchy_elapsed_ms = elapsed_ms_since(hierarchy_start_time); - RAFT_LOG_INFO("Layered HNSW: hierarchy levels generated in %ld ms (max_level=%d, promoted=%zu)", - hierarchy_elapsed_ms, - hierarchy.max_level(), - hierarchy.promoted_count()); + RAFT_LOG_INFO( + "Layered HNSW artifact: hierarchy levels generated in %ld ms (max_level=%d promoted=%zu)", + hierarchy_elapsed_ms, + hierarchy.max_level(), + hierarchy.promoted_count()); layered_hnsw_file_metadata metadata; metadata.n_rows = n_rows; @@ -1015,12 +1054,12 @@ auto serialize_to_layered_hnsw_from_disk( cuvs::util::write_large_file( artifact_fd, hierarchy.levels.data(), metadata.levels_bytes, levels_offset); const auto levels_elapsed_ms = elapsed_ms_since(levels_start_time); - RAFT_LOG_INFO("Layered HNSW: wrote levels section in %ld ms (%.2f GiB, %.2f GiB/s)", + RAFT_LOG_INFO("Layered HNSW artifact: levels section written in %ld ms (%.2f GiB, %.2f GiB/s)", levels_elapsed_ms, to_gib(metadata.levels_bytes), throughput_gib_per_s(metadata.levels_bytes, levels_elapsed_ms)); - RAFT_LOG_INFO("Layered HNSW: writing hnswlib-ready base links"); + RAFT_LOG_INFO("Layered HNSW artifact: writing hnswlib-ready base topology section"); const auto layer0_start_time = std::chrono::steady_clock::now(); write_layered_base_links_from_disk(index_, artifact_fd, @@ -1029,14 +1068,16 @@ auto serialize_to_layered_hnsw_from_disk( metadata.base_link_row_bytes, metadata.maxM0); const auto layer0_elapsed_ms = elapsed_ms_since(layer0_start_time); - RAFT_LOG_INFO("Layered HNSW: base links written in %ld ms (%.2f GiB, %.2f GiB/s effective)", - layer0_elapsed_ms, - to_gib(metadata.base_links_bytes), - throughput_gib_per_s(metadata.base_links_bytes, layer0_elapsed_ms)); + static_cast(layer0_elapsed_ms); + RAFT_LOG_INFO( + "Layered HNSW artifact: base topology section written in %ld ms (%.2f GiB, %.2f GiB/s)", + layer0_elapsed_ms, + to_gib(metadata.base_nodes_bytes + metadata.base_links_bytes), + throughput_gib_per_s(metadata.base_nodes_bytes + metadata.base_links_bytes, layer0_elapsed_ms)); size_t upper_graph_bytes_written = 0; if (hierarchy.hist.size() > 1) { - RAFT_LOG_INFO("Layered HNSW: gathering promoted vectors"); + RAFT_LOG_INFO("Layered HNSW artifact: gathering promoted vectors"); const auto gather_start_time = std::chrono::steady_clock::now(); auto host_query_set = raft::make_host_matrix(static_cast(hierarchy.promoted_count()), dim); @@ -1050,10 +1091,11 @@ auto serialize_to_layered_hnsw_from_disk( } const auto gather_elapsed_ms = elapsed_ms_since(gather_start_time); const auto gathered_bytes = hierarchy.promoted_count() * dim * sizeof(T); - RAFT_LOG_INFO("Layered HNSW: gathered promoted vectors in %ld ms (%.2f GiB copied, %.2f GiB/s)", - gather_elapsed_ms, - to_gib(gathered_bytes), - throughput_gib_per_s(gathered_bytes, gather_elapsed_ms)); + RAFT_LOG_INFO( + "Layered HNSW artifact: promoted vectors gathered in %ld ms (%.2f GiB copied, %.2f GiB/s)", + gather_elapsed_ms, + to_gib(gathered_bytes), + throughput_gib_per_s(gathered_bytes, gather_elapsed_ms)); auto promoted_dataset = raft::make_host_matrix_view( host_query_set.data_handle(), host_query_set.extent(0), host_query_set.extent(1)); @@ -1066,10 +1108,10 @@ auto serialize_to_layered_hnsw_from_disk( index_.metric(), [&](size_t pt_level, size_t start_idx, auto& host_neighbors) { const auto& layer = metadata.layers[pt_level - 1]; - RAFT_LOG_INFO("Layered HNSW: writing upper layer %zu (%zu rows, degree %zu)", - layer.level, - layer.row_count, - layer.degree); + RAFT_LOG_DEBUG("Layered HNSW artifact: writing upper layer level=%zu rows=%zu degree=%zu", + layer.level, + layer.row_count, + layer.degree); const auto layer_write_start_time = std::chrono::steady_clock::now(); const size_t target_batch_bytes = 64 * 1024 * 1024; const size_t row_bytes = sizeof(IdxT) + metadata.upper_link_row_bytes; @@ -1112,28 +1154,29 @@ auto serialize_to_layered_hnsw_from_disk( const auto layer_bytes = layer.row_count * (sizeof(IdxT) + metadata.upper_link_row_bytes); upper_graph_bytes_written += layer_bytes; const auto layer_write_elapsed_ms = elapsed_ms_since(layer_write_start_time); - RAFT_LOG_INFO("Layered HNSW: upper layer %zu written in %ld ms (%.2f GiB, %.2f GiB/s)", - layer.level, - layer_write_elapsed_ms, - to_gib(layer_bytes), - throughput_gib_per_s(layer_bytes, layer_write_elapsed_ms)); + static_cast(layer_write_elapsed_ms); + RAFT_LOG_INFO( + "Layered HNSW artifact: upper layer level=%zu written in %ld ms (%.2f GiB, %.2f GiB/s)", + layer.level, + layer_write_elapsed_ms, + to_gib(layer_bytes), + throughput_gib_per_s(layer_bytes, layer_write_elapsed_ms)); }); const auto upper_layers_elapsed_ms = elapsed_ms_since(upper_layers_start_time); RAFT_LOG_INFO( - "Layered HNSW: upper layers generated and written in %ld ms (%.2f GiB written, %.2f " - "GiB/s effective)", + "Layered HNSW artifact: upper layers generated and written in %ld ms (%.2f GiB, %.2f " + "GiB/s)", upper_layers_elapsed_ms, to_gib(upper_graph_bytes_written), throughput_gib_per_s(upper_graph_bytes_written, upper_layers_elapsed_ms)); } const auto total_elapsed_ms = elapsed_ms_since(total_start_time); - RAFT_LOG_INFO( - "Layered HNSW index written to: %s in %ld ms (artifact %.2f GiB, %.2f GiB/s effective)", - artifact_file.string().c_str(), - total_elapsed_ms, - to_gib(final_file_size), - throughput_gib_per_s(final_file_size, total_elapsed_ms)); + RAFT_LOG_INFO("Layered HNSW artifact: wrote %s in %ld ms (artifact=%.2f GiB, %.2f GiB/s)", + artifact_file.string().c_str(), + total_elapsed_ms, + to_gib(final_file_size), + throughput_gib_per_s(final_file_size, total_elapsed_ms)); return artifact_file.string(); } @@ -2096,7 +2139,13 @@ auto deserialize_layered_hnsw(raft::resources const& res, RAFT_EXPECTS(dataset_file.shape.size() == 2 && dataset_file.shape[0] == metadata.n_rows && dataset_file.shape[1] == metadata.dim, - "Layered HNSW dataset shape mismatch"); + "Layered HNSW dataset shape mismatch: artifact rows=%zu dim=%zu, dataset rows=%zu " + "dim=%zu path=%s", + metadata.n_rows, + metadata.dim, + dataset_file.shape.size() > 0 ? dataset_file.shape[0] : 0, + dataset_file.shape.size() > 1 ? dataset_file.shape[1] : 0, + params.dataset_path.c_str()); RAFT_EXPECTS(metadata.levels_bytes == metadata.n_rows * sizeof(uint8_t), "Layered HNSW levels section size mismatch"); RAFT_EXPECTS(metadata.base_nodes_bytes == metadata.n_rows * sizeof(uint32_t), @@ -2117,9 +2166,32 @@ auto deserialize_layered_hnsw(raft::resources const& res, dataset_open_elapsed_ms, params.dataset_path.c_str()); + const auto dataset_total_bytes = metadata.n_rows * metadata.dim * sizeof(T); + const auto deserialize_progress_total_bytes = + metadata.levels_bytes + dataset_total_bytes + metadata.base_nodes_bytes + + metadata.base_links_bytes + metadata.upper_nodes_bytes + metadata.upper_links_bytes; + size_t deserialize_progress_bytes = 0; + size_t deserialize_next_progress_percent = 10; + auto log_deserialize_progress = [&](size_t bytes_loaded) { + if (deserialize_progress_total_bytes == 0) { return; } + deserialize_progress_bytes = + std::min(deserialize_progress_total_bytes, deserialize_progress_bytes + bytes_loaded); + const auto progress_percent = + (deserialize_progress_bytes * 100) / deserialize_progress_total_bytes; + const auto progress_step = std::min(100, (progress_percent / 10) * 10); + if (progress_step >= deserialize_next_progress_percent) { + RAFT_LOG_INFO("Layered HNSW load: deserialize progress %zu%% (%.2f/%.2f GiB)", + progress_step, + to_gib(deserialize_progress_bytes), + to_gib(deserialize_progress_total_bytes)); + deserialize_next_progress_percent = progress_step + 10; + } + }; + const auto levels_start_time = std::chrono::steady_clock::now(); std::vector levels_u8(metadata.n_rows); cuvs::util::read_large_file(artifact_fd, levels_u8.data(), metadata.levels_bytes, levels_offset); + log_deserialize_progress(metadata.levels_bytes); const auto max_level_in_levels = *std::max_element(levels_u8.begin(), levels_u8.end()); RAFT_EXPECTS(static_cast(max_level_in_levels) == metadata.maxlevel, "Layered HNSW levels max level (%d) does not match artifact maxlevel (%d)", @@ -2203,6 +2275,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, throw std::runtime_error("Not enough memory to allocate HNSW upper linklists"); } base_copy_time += std::chrono::steady_clock::now() - batch_timer; + log_deserialize_progress(current_dataset_bytes); } const auto base_elapsed_ms = elapsed_ms_since(base_start_time); RAFT_LOG_INFO( @@ -2257,6 +2330,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, } RAFT_EXPECTS(!invalid_node_id, "Invalid base-layer node id in layered HNSW artifact"); base_topology_copy_time += std::chrono::steady_clock::now() - batch_timer; + log_deserialize_progress(current_base_topology_bytes); } const auto base_topology_elapsed_ms = elapsed_ms_since(base_topology_start_time); RAFT_LOG_INFO( @@ -2320,6 +2394,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, RAFT_EXPECTS(!invalid_node_level, "Layered HNSW artifact references a node at an invalid upper level"); upper_copy_time += std::chrono::steady_clock::now() - batch_timer; + log_deserialize_progress(current_upper_bytes); } } const auto upper_elapsed_ms = elapsed_ms_since(upper_start_time); @@ -2412,15 +2487,14 @@ std::unique_ptr> build(raft::resources const& res, cagra_ace_params.max_gpu_memory_gb = ace_params.max_gpu_memory_gb; cagra_params.graph_build_params = cagra_ace_params; - RAFT_LOG_INFO( - "hnsw::build - Building HNSW index using ACE with %zu partitions, ef_construction=%zu", - ace_params.npartitions, - ace_params.ef_construction); + RAFT_LOG_INFO("HNSW ACE build: building CAGRA graph (partitions=%zu ef_construction=%zu)", + cagra_ace_params.npartitions, + cagra_ace_params.ef_construction); // Build CAGRA index using ACE auto cagra_index = cuvs::neighbors::cagra::build(res, cagra_params, dataset); - RAFT_LOG_INFO("hnsw::build - Converting CAGRA index to HNSW format"); + RAFT_LOG_INFO("HNSW ACE build: converting CAGRA graph to HNSW hierarchy"); // Convert CAGRA index to HNSW index return from_cagra(res, params, cagra_index, dataset); From 7a761cf90a9cabee08ddd4bdb359bbf66056ac3c Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 9 Jun 2026 09:51:11 +0200 Subject: [PATCH 07/12] Address review feedback --- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu | 3 +- .../ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h | 28 ++++++++++--------- cpp/src/neighbors/detail/hnsw.hpp | 6 ++++ 3 files changed, 22 insertions(+), 15 deletions(-) diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu index d4791dde11..65adc926f8 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu @@ -28,8 +28,7 @@ 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") { + } else if (conf.at("hierarchy") == "gpu_layered_on_disk") { hnsw_params.hierarchy = cuvs::neighbors::hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK; } else { THROW("Invalid value for hierarchy: %s", conf.at("hierarchy").get().c_str()); diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index 0104ccdb3e..8eb1b72991 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h @@ -14,7 +14,7 @@ namespace cuvs::bench { -inline void copy_file_overwrite(const std::filesystem::path& src, const std::filesystem::path& dst) +inline void move_file_overwrite(const std::filesystem::path& src, const std::filesystem::path& dst) { std::error_code ec; if (src == dst || @@ -22,11 +22,18 @@ inline void copy_file_overwrite(const std::filesystem::path& src, const std::fil 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()); + if (std::filesystem::exists(dst, ec)) { std::filesystem::remove(dst, ec); } + + std::filesystem::rename(src, dst, ec); + if (ec) { + // Rename fails across filesystems. Fall back to copy followed by removal of the source. + ec.clear(); + 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 move '%s' to '%s'.", src_str.c_str(), dst_str.c_str()); + std::filesystem::remove(src, ec); + } } template @@ -154,7 +161,7 @@ void cuvs_cagra_hnswlib::save(const std::string& file) const "Layered HNSW artifact '%s' does not exist.", src_artifact.c_str()); - copy_file_overwrite(src_artifact, std::filesystem::path(file)); + move_file_overwrite(src_artifact, std::filesystem::path(file)); return; } @@ -164,12 +171,7 @@ void cuvs_cagra_hnswlib::save(const std::string& file) const RAFT_EXPECTS(std::filesystem::exists(index_filename), "Index file '%s' does not exist.", index_filename.c_str()); - if (std::filesystem::exists(file)) { std::filesystem::remove(file); } - // might fail when using 2 different filesystems - std::error_code ec; - std::filesystem::rename(index_filename, file, ec); - RAFT_EXPECTS( - !ec, "Failed to rename index file '%s' to '%s'.", index_filename.c_str(), file.c_str()); + move_file_overwrite(std::filesystem::path(index_filename), std::filesystem::path(file)); } else { cuvs::neighbors::hnsw::serialize(handle_, file, *(hnsw_index_.get())); } diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 5be875d804..28d4e7db56 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -878,6 +878,9 @@ void write_layered_base_links_from_disk(const cuvs::neighbors::cagra::index(current_batch_size); ++batch_idx) { @@ -1123,6 +1126,9 @@ auto serialize_to_layered_hnsw_from_disk( std::fill(link_buffer.begin(), link_buffer.begin() + current_batch_size * metadata.upper_link_row_bytes, 0); + // Upper-layer rows stay in promoted (per-level) order for sequential writes. node_buffer + // records each node's original ID and neighbors are remapped to original IDs. The search + // node scatters each row into get_linklist(original_id, level) on load. #pragma omp parallel for for (int64_t batch_idx = 0; batch_idx < static_cast(current_batch_size); ++batch_idx) { From ea1a96c07f63567e248234c3d02527209ea16f15 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 9 Jun 2026 09:53:41 +0200 Subject: [PATCH 08/12] Replace JSON header with binary header --- cpp/src/neighbors/detail/hnsw.hpp | 308 +++++++++++++++--------------- 1 file changed, 158 insertions(+), 150 deletions(-) diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 28d4e7db56..41fe241f98 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -38,7 +38,6 @@ #include #include #include -#include #include #include #include @@ -539,13 +538,55 @@ struct layered_hnsw_file_metadata { std::vector layers; }; +enum class layered_hnsw_dtype : uint32_t { + unknown = 0, + float32 = 1, + float16 = 2, + uint8 = 3, + int8 = 4, +}; + +// The layered HNSW artifact begins with this fixed-size POD header, immediately followed by +// `num_layers` layered_hnsw_layer_descriptor records and then the payload sections. struct layered_hnsw_file_header { char magic[32]; + uint64_t n_rows; + uint64_t dim; + uint64_t M; + uint64_t maxM; + uint64_t maxM0; + uint64_t ef_construction; + uint64_t base_degree; + uint64_t levels_bytes; + uint64_t base_nodes_bytes; + uint64_t base_link_row_bytes; + uint64_t base_links_bytes; + uint64_t upper_nodes_count; + uint64_t upper_nodes_bytes; + uint64_t upper_link_row_bytes; + uint64_t upper_links_bytes; + double mult; uint32_t version; - uint32_t reserved; - uint64_t metadata_size; - uint64_t metadata_offset; + uint32_t dtype; + uint32_t metric; + uint32_t num_layers; + int32_t maxlevel; + int32_t enterpoint_node; + uint32_t reserved0; + uint32_t reserved1; +}; +static_assert(sizeof(layered_hnsw_file_header) == 192, + "layered_hnsw_file_header must keep a fixed 192-byte on-disk layout"); + +struct layered_hnsw_layer_descriptor { + uint64_t level; + uint64_t row_count; + uint64_t degree; + uint64_t node_offset; + uint64_t link_offset; }; +static_assert(sizeof(layered_hnsw_layer_descriptor) == 40, + "layered_hnsw_layer_descriptor must keep a fixed 40-byte on-disk layout"); constexpr const char* layered_hnsw_magic = "CUVS_HNSW_LAYERED"; constexpr uint32_t layered_hnsw_version = 1; @@ -557,18 +598,29 @@ inline auto align_up(size_t value, size_t alignment) -> size_t } template -auto layered_dtype_name() -> const char* +constexpr auto layered_dtype_code() -> layered_hnsw_dtype { if constexpr (std::is_same_v) { - return "float32"; + return layered_hnsw_dtype::float32; } else if constexpr (std::is_same_v) { - return "float16"; + return layered_hnsw_dtype::float16; } else if constexpr (std::is_same_v) { - return "uint8"; + return layered_hnsw_dtype::uint8; } else if constexpr (std::is_same_v) { - return "int8"; + return layered_hnsw_dtype::int8; } else { - return "unknown"; + return layered_hnsw_dtype::unknown; + } +} + +inline auto layered_dtype_name(layered_hnsw_dtype dtype) -> const char* +{ + switch (dtype) { + case layered_hnsw_dtype::float32: return "float32"; + case layered_hnsw_dtype::float16: return "float16"; + case layered_hnsw_dtype::uint8: return "uint8"; + case layered_hnsw_dtype::int8: return "int8"; + default: return "unknown"; } } @@ -582,128 +634,58 @@ inline auto metric_name(cuvs::distance::DistanceType metric) -> const char* } template -auto make_layered_hnsw_metadata_json(const layered_hnsw_file_metadata& metadata, - cuvs::distance::DistanceType metric) -> std::string -{ - std::stringstream os; - os << "{\n"; - os << " \"format\": \"cuvs_hnsw_layered\",\n"; - os << " \"version\": " << layered_hnsw_version << ",\n"; - os << " \"n_rows\": " << metadata.n_rows << ",\n"; - os << " \"dim\": " << metadata.dim << ",\n"; - os << " \"dtype\": \"" << layered_dtype_name() << "\",\n"; - os << " \"metric\": \"" << metric_name(metric) << "\",\n"; - os << " \"row_order\": \"original_id\",\n"; - os << " \"base_link_order\": \"base_nodes\",\n"; - os << " \"M\": " << metadata.M << ",\n"; - os << " \"maxM\": " << metadata.maxM << ",\n"; - os << " \"maxM0\": " << metadata.maxM0 << ",\n"; - os << " \"ef_construction\": " << metadata.ef_construction << ",\n"; - os << " \"mult\": " << metadata.mult << ",\n"; - os << " \"maxlevel\": " << metadata.maxlevel << ",\n"; - os << " \"enterpoint_node\": " << metadata.enterpoint_node << ",\n"; - os << " \"base_degree\": " << metadata.base_degree << ",\n"; - os << " \"levels_bytes\": " << metadata.levels_bytes << ",\n"; - os << " \"base_nodes_bytes\": " << metadata.base_nodes_bytes << ",\n"; - os << " \"base_link_row_bytes\": " << metadata.base_link_row_bytes << ",\n"; - os << " \"base_links_bytes\": " << metadata.base_links_bytes << ",\n"; - os << " \"upper_nodes_count\": " << metadata.upper_nodes_count << ",\n"; - os << " \"upper_nodes_bytes\": " << metadata.upper_nodes_bytes << ",\n"; - os << " \"upper_link_row_bytes\": " << metadata.upper_link_row_bytes << ",\n"; - os << " \"upper_links_bytes\": " << metadata.upper_links_bytes << ",\n"; - os << " \"layers\": [\n"; - for (size_t i = 0; i < metadata.layers.size(); ++i) { - const auto& layer = metadata.layers[i]; - os << " {\"level\": " << layer.level << ", \"row_count\": " << layer.row_count - << ", \"degree\": " << layer.degree << ", \"node_offset\": " << layer.node_offset - << ", \"link_offset\": " << layer.link_offset << "}"; - os << (i + 1 == metadata.layers.size() ? "\n" : ",\n"); - } - os << " ]\n"; - os << "}\n"; - return os.str(); -} - -inline auto json_find_key(const std::string& json, const std::string& key, size_t start = 0) - -> size_t +auto make_layered_hnsw_header(const layered_hnsw_file_metadata& metadata, + cuvs::distance::DistanceType metric) -> layered_hnsw_file_header { - const auto pos = json.find("\"" + key + "\"", start); - RAFT_EXPECTS( - pos != std::string::npos, "Cannot find key '%s' in layered HNSW metadata", key.c_str()); - return pos; -} - -inline auto json_parse_size(const std::string& json, const std::string& key) -> size_t -{ - auto pos = json_find_key(json, key); - auto colon = json.find(':', pos); - RAFT_EXPECTS(colon != std::string::npos, "Malformed JSON near key '%s'", key.c_str()); - auto begin = json.find_first_of("0123456789", colon + 1); - auto end = json.find_first_not_of("0123456789", begin); - RAFT_EXPECTS( - begin != std::string::npos, "Malformed integer JSON value for key '%s'", key.c_str()); - return static_cast(std::stoull(json.substr(begin, end - begin))); -} - -inline auto json_parse_double(const std::string& json, const std::string& key) -> double -{ - auto pos = json_find_key(json, key); - auto colon = json.find(':', pos); - RAFT_EXPECTS(colon != std::string::npos, "Malformed JSON near key '%s'", key.c_str()); - auto begin = json.find_first_of("0123456789-.", colon + 1); - auto end = json.find_first_not_of("0123456789-.eE+", begin); - RAFT_EXPECTS(begin != std::string::npos, "Malformed double JSON value for key '%s'", key.c_str()); - return std::stod(json.substr(begin, end - begin)); -} - -inline auto json_parse_layer_field(const std::string& object, const std::string& key) -> size_t -{ - return json_parse_size(object, key); + layered_hnsw_file_header header{}; + std::strncpy(header.magic, layered_hnsw_magic, sizeof(header.magic) - 1); + header.version = layered_hnsw_version; + header.dtype = static_cast(layered_dtype_code()); + header.metric = static_cast(metric); + header.num_layers = static_cast(metadata.layers.size()); + header.n_rows = metadata.n_rows; + header.dim = metadata.dim; + header.M = metadata.M; + header.maxM = metadata.maxM; + header.maxM0 = metadata.maxM0; + header.ef_construction = metadata.ef_construction; + header.base_degree = metadata.base_degree; + header.levels_bytes = metadata.levels_bytes; + header.base_nodes_bytes = metadata.base_nodes_bytes; + header.base_link_row_bytes = metadata.base_link_row_bytes; + header.base_links_bytes = metadata.base_links_bytes; + header.upper_nodes_count = metadata.upper_nodes_count; + header.upper_nodes_bytes = metadata.upper_nodes_bytes; + header.upper_link_row_bytes = metadata.upper_link_row_bytes; + header.upper_links_bytes = metadata.upper_links_bytes; + header.mult = metadata.mult; + header.maxlevel = metadata.maxlevel; + header.enterpoint_node = metadata.enterpoint_node; + return header; } -inline auto parse_layered_hnsw_metadata(const std::string& json) -> layered_hnsw_file_metadata +inline auto layered_hnsw_metadata_from_header(const layered_hnsw_file_header& header) + -> layered_hnsw_file_metadata { layered_hnsw_file_metadata metadata; - metadata.n_rows = json_parse_size(json, "n_rows"); - metadata.dim = json_parse_size(json, "dim"); - metadata.M = json_parse_size(json, "M"); - metadata.maxM = json_parse_size(json, "maxM"); - metadata.maxM0 = json_parse_size(json, "maxM0"); - metadata.ef_construction = json_parse_size(json, "ef_construction"); - metadata.mult = json_parse_double(json, "mult"); - metadata.maxlevel = static_cast(json_parse_size(json, "maxlevel")); - metadata.enterpoint_node = static_cast(json_parse_size(json, "enterpoint_node")); - metadata.base_degree = json_parse_size(json, "base_degree"); - metadata.levels_bytes = json_parse_size(json, "levels_bytes"); - metadata.base_nodes_bytes = json_parse_size(json, "base_nodes_bytes"); - metadata.base_link_row_bytes = json_parse_size(json, "base_link_row_bytes"); - metadata.base_links_bytes = json_parse_size(json, "base_links_bytes"); - metadata.upper_nodes_count = json_parse_size(json, "upper_nodes_count"); - metadata.upper_nodes_bytes = json_parse_size(json, "upper_nodes_bytes"); - metadata.upper_link_row_bytes = json_parse_size(json, "upper_link_row_bytes"); - metadata.upper_links_bytes = json_parse_size(json, "upper_links_bytes"); - - auto layers_pos = json_find_key(json, "layers"); - auto array_open = json.find('[', layers_pos); - auto array_end = json.find(']', array_open); - RAFT_EXPECTS(array_open != std::string::npos && array_end != std::string::npos, - "Malformed layers array in layered HNSW metadata"); - - size_t pos = array_open; - while (true) { - auto object_open = json.find('{', pos); - if (object_open == std::string::npos || object_open > array_end) { break; } - auto object_end = json.find('}', object_open); - RAFT_EXPECTS(object_end != std::string::npos && object_end < array_end, - "Malformed layer object in layered HNSW metadata"); - const auto object = json.substr(object_open, object_end - object_open + 1); - metadata.layers.push_back({json_parse_layer_field(object, "level"), - json_parse_layer_field(object, "row_count"), - json_parse_layer_field(object, "degree"), - json_parse_layer_field(object, "node_offset"), - json_parse_layer_field(object, "link_offset")}); - pos = object_end + 1; - } + metadata.n_rows = header.n_rows; + metadata.dim = header.dim; + metadata.M = header.M; + metadata.maxM = header.maxM; + metadata.maxM0 = header.maxM0; + metadata.ef_construction = header.ef_construction; + metadata.base_degree = header.base_degree; + metadata.levels_bytes = header.levels_bytes; + metadata.base_nodes_bytes = header.base_nodes_bytes; + metadata.base_link_row_bytes = header.base_link_row_bytes; + metadata.base_links_bytes = header.base_links_bytes; + metadata.upper_nodes_count = header.upper_nodes_count; + metadata.upper_nodes_bytes = header.upper_nodes_bytes; + metadata.upper_link_row_bytes = header.upper_link_row_bytes; + metadata.upper_links_bytes = header.upper_links_bytes; + metadata.mult = header.mult; + metadata.maxlevel = header.maxlevel; + metadata.enterpoint_node = header.enterpoint_node; return metadata; } @@ -1027,15 +1009,21 @@ auto serialize_to_layered_hnsw_from_disk( metadata.upper_nodes_bytes = metadata.upper_nodes_count * sizeof(IdxT); metadata.upper_links_bytes = next_upper_link_offset * metadata.upper_link_row_bytes; - const auto metadata_json = make_layered_hnsw_metadata_json(metadata, index_.metric()); - layered_hnsw_file_header header{}; - std::strncpy(header.magic, layered_hnsw_magic, sizeof(header.magic) - 1); - header.version = layered_hnsw_version; - header.metadata_size = metadata_json.size(); - header.metadata_offset = sizeof(layered_hnsw_file_header); + const auto header = make_layered_hnsw_header(metadata, index_.metric()); + std::vector layer_descriptors(metadata.layers.size()); + for (size_t i = 0; i < metadata.layers.size(); ++i) { + const auto& layer = metadata.layers[i]; + layer_descriptors[i] = {static_cast(layer.level), + static_cast(layer.row_count), + static_cast(layer.degree), + static_cast(layer.node_offset), + static_cast(layer.link_offset)}; + } + const auto descriptors_offset = sizeof(layered_hnsw_file_header); + const auto descriptors_bytes = layer_descriptors.size() * sizeof(layered_hnsw_layer_descriptor); const auto payload_offset = - align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); + align_up(descriptors_offset + descriptors_bytes, layered_hnsw_alignment); const auto levels_offset = payload_offset; const auto base_nodes_offset = levels_offset + metadata.levels_bytes; const auto base_links_offset = base_nodes_offset + metadata.base_nodes_bytes; @@ -1050,8 +1038,10 @@ auto serialize_to_layered_hnsw_from_disk( artifact_file.string().c_str(), std::strerror(fallocate_result)); cuvs::util::write_large_file(artifact_fd, &header, sizeof(header), 0); - cuvs::util::write_large_file( - artifact_fd, metadata_json.data(), metadata_json.size(), header.metadata_offset); + if (descriptors_bytes > 0) { + cuvs::util::write_large_file( + artifact_fd, layer_descriptors.data(), descriptors_bytes, descriptors_offset); + } const auto levels_start_time = std::chrono::steady_clock::now(); cuvs::util::write_large_file( @@ -2089,22 +2079,40 @@ auto deserialize_layered_hnsw(raft::resources const& res, RAFT_EXPECTS(header.version == layered_hnsw_version, "Unsupported layered HNSW artifact version %u", header.version); - RAFT_EXPECTS(header.metadata_offset >= sizeof(layered_hnsw_file_header), - "Invalid layered HNSW metadata offset"); - RAFT_EXPECTS(header.metadata_size > 0, "Invalid layered HNSW metadata size"); - const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); - RAFT_EXPECTS(header.metadata_offset <= artifact_size && - header.metadata_size <= artifact_size - header.metadata_offset, - "Layered HNSW metadata range is outside artifact: offset=%zu size=%zu " + RAFT_EXPECTS(header.dtype == static_cast(layered_dtype_code()), + "Layered HNSW artifact dtype (%s) does not match requested dtype (%s)", + layered_dtype_name(static_cast(header.dtype)), + layered_dtype_name(layered_dtype_code())); + RAFT_EXPECTS(header.metric == static_cast(metric), + "Layered HNSW artifact metric (%s) does not match requested metric (%s)", + metric_name(static_cast(header.metric)), + metric_name(metric)); + + const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); + const auto descriptors_offset = sizeof(layered_hnsw_file_header); + const auto descriptors_bytes = + static_cast(header.num_layers) * sizeof(layered_hnsw_layer_descriptor); + RAFT_EXPECTS(descriptors_offset + descriptors_bytes <= artifact_size, + "Layered HNSW layer descriptors are outside artifact: offset=%zu size=%zu " "artifact=%zu", - static_cast(header.metadata_offset), - static_cast(header.metadata_size), + descriptors_offset, + descriptors_bytes, artifact_size); - std::string metadata_json(header.metadata_size, '\0'); - cuvs::util::read_large_file( - artifact_fd, metadata_json.data(), metadata_json.size(), header.metadata_offset); - const auto metadata = parse_layered_hnsw_metadata(metadata_json); + std::vector layer_descriptors(header.num_layers); + if (descriptors_bytes > 0) { + cuvs::util::read_large_file( + artifact_fd, layer_descriptors.data(), descriptors_bytes, descriptors_offset); + } + auto metadata = layered_hnsw_metadata_from_header(header); + metadata.layers.reserve(layer_descriptors.size()); + for (const auto& descriptor : layer_descriptors) { + metadata.layers.push_back({static_cast(descriptor.level), + static_cast(descriptor.row_count), + static_cast(descriptor.degree), + static_cast(descriptor.node_offset), + static_cast(descriptor.link_offset)}); + } const auto metadata_elapsed_ms = elapsed_ms_since(metadata_start_time); RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); @@ -2122,7 +2130,7 @@ auto deserialize_layered_hnsw(raft::resources const& res, "Layered HNSW deserialization requires index_params.dataset_path"); const auto payload_offset = - align_up(header.metadata_offset + header.metadata_size, layered_hnsw_alignment); + align_up(descriptors_offset + descriptors_bytes, layered_hnsw_alignment); const auto levels_offset = payload_offset; const auto base_nodes_offset = levels_offset + metadata.levels_bytes; const auto base_links_offset = base_nodes_offset + metadata.base_nodes_bytes; From ad3848a358cd27cac498f08e285468b18bc04652 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Fri, 12 Jun 2026 15:01:18 +0200 Subject: [PATCH 09/12] Add GPU_LAYERED_ON_DISK support for HNSW materialization - `materialize_to_hnswlib` enables disk-to-disk materialization of a layered HNSW artifact to a standard hnswlib index fiel. - Adds bindings --- c/include/cuvs/neighbors/hnsw.h | 86 +- c/src/neighbors/hnsw.cpp | 43 +- cpp/include/cuvs/neighbors/hnsw.hpp | 69 ++ cpp/include/cuvs/util/file_io.hpp | 13 + cpp/src/neighbors/detail/hnsw.hpp | 891 ++++++++++++++++-- cpp/src/neighbors/hnsw.cpp | 13 + cpp/src/util/file_io.cpp | 21 +- cpp/tests/neighbors/ann_hnsw_ace.cuh | 202 ++++ .../ann_hnsw_ace/test_float_uint32_t.cu | 10 + .../ann_hnsw_ace/test_half_uint32_t.cu | 10 + .../ann_hnsw_ace/test_int8_t_uint32_t.cu | 10 + .../ann_hnsw_ace/test_uint8_t_uint32_t.cu | 10 + examples/cpp/src/hnsw_ace_layered_example.cu | 122 ++- .../main/java/com/nvidia/cuvs/HnswIndex.java | 38 + .../java/com/nvidia/cuvs/HnswIndexParams.java | 54 +- .../nvidia/cuvs/HnswMaterializeParams.java | 134 +++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 22 + .../cuvs/internal/CuVSParamsHelper.java | 21 + .../nvidia/cuvs/internal/HnswIndexImpl.java | 76 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 13 + python/cuvs/cuvs/neighbors/hnsw/__init__.py | 4 + python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd | 23 + python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 177 +++- 23 files changed, 1943 insertions(+), 119 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswMaterializeParams.java diff --git a/c/include/cuvs/neighbors/hnsw.h b/c/include/cuvs/neighbors/hnsw.h index 15eb1b0569..9f4c1c4445 100644 --- a/c/include/cuvs/neighbors/hnsw.h +++ b/c/include/cuvs/neighbors/hnsw.h @@ -35,7 +35,9 @@ enum cuvsHnswHierarchy { /* Full hierarchy is built using the CPU */ CPU = 1, /* Full hierarchy is built using the GPU */ - GPU = 2 + GPU = 2, + /* GPU-built hierarchy stored as a layered on-disk topology artifact */ + GPU_LAYERED_ON_DISK = 3 }; /** @@ -130,6 +132,14 @@ struct cuvsHnswIndexParams { * Set to nullptr for default behavior (from_cagra conversion). */ cuvsHnswAceParams_t ace_params; + /** + * Local dataset path used by layered HNSW deserialization. + * + * When `hierarchy == GPU_LAYERED_ON_DISK`, the index artifact stores graph topology only. + * `cuvsHnswDeserialize` loads the original-ID-ordered vectors from this local dataset path to + * reconstruct an in-memory HNSW index. Set to nullptr (default) for all other hierarchies. + */ + const char* dataset_path; }; typedef struct cuvsHnswIndexParams* cuvsHnswIndexParams_t; @@ -601,6 +611,80 @@ CUVS_EXPORT cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, * @} */ +/** + * @defgroup hnsw_c_index_materialize Materialize a layered HNSW artifact to an hnswlib index + * @{ + */ + +/** + * @brief Parameters for materializing a layered HNSW artifact into an hnswlib index on disk. + */ +struct cuvsHnswMaterializeParams { + /** + * Local dataset path holding the original-ID-ordered vectors used to build the artifact. + * + * Supported formats match layered deserialization: `.npy` and ANN benchmark `*.bin` files with a + * `[uint32 rows, uint32 cols]` header (`.fbin`, `.f16bin`, `.u8bin`, `.i8bin`). + */ + const char* dataset_path; + /** + * Upper bound on host memory (in GiB) used for the base-topology reorder buffer. + * + * When `<= 0`, the whole base topology is reordered in a single in-memory pass (no temporary + * files). When set, the base topology is reordered through bucketed temporary files so that peak + * host memory stays close to this budget. + */ + double max_host_memory_gb; + /** Number of host threads to use. When `0`, the maximum number of threads is used. */ + int num_threads; +}; + +typedef struct cuvsHnswMaterializeParams* cuvsHnswMaterializeParams_t; + +/** + * @brief Allocate HNSW materialize params, and populate with default values + * + * @param[in] params cuvsHnswMaterializeParams_t to allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsHnswMaterializeParamsCreate(cuvsHnswMaterializeParams_t* params); + +/** + * @brief De-allocate HNSW materialize params + * + * @param[in] params cuvsHnswMaterializeParams_t to de-allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsHnswMaterializeParamsDestroy(cuvsHnswMaterializeParams_t params); + +/** + * @brief Materialize a layered HNSW artifact into a standard hnswlib index file on disk. + * + * Materializes a `GPU_LAYERED_ON_DISK` artifact (graph topology only, stored in ACE order) plus a + * local dataset into a standard hnswlib index file, without ever holding the full materialized + * index in host memory. The resulting file is compatible with the original hnswlib library and can + * be read back through `cuvsHnswDeserialize` with `hierarchy == CPU`. The element data type + * (`float`, `half`, `uint8_t` or `int8_t`) is read from the artifact header. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params cuvsHnswMaterializeParams_t materialization parameters + * @param[in] layered_artifact_path path to the layered HNSW artifact + * @param[in] output_path path to the hnswlib index file to write + * @param[in] dim the dimension of the vectors in the index + * @param[in] metric the distance metric used to build the index + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsHnswMaterializeToHnswlib(cuvsResources_t res, + cuvsHnswMaterializeParams_t params, + const char* layered_artifact_path, + const char* output_path, + int dim, + cuvsDistanceType metric); + +/** + * @} + */ + #ifdef __cplusplus } #endif diff --git a/c/src/neighbors/hnsw.cpp b/c/src/neighbors/hnsw.cpp index c69eda0ca0..c791c004a3 100644 --- a/c/src/neighbors/hnsw.cpp +++ b/c/src/neighbors/hnsw.cpp @@ -143,6 +143,8 @@ void* _deserialize(cuvsResources_t res, cuvs::neighbors::hnsw::index* index = nullptr; auto cpp_params = cuvs::neighbors::hnsw::index_params(); cpp_params.hierarchy = static_cast(params->hierarchy); + // Required by GPU_LAYERED_ON_DISK deserialization to locate the original-ID-ordered vectors. + if (params->dataset_path != nullptr) { cpp_params.dataset_path = std::string(params->dataset_path); } auto metric_type = static_cast(metric); cuvs::neighbors::hnsw::deserialize( *res_ptr, cpp_params, std::string(filename), dim, metric_type, &index); @@ -174,7 +176,8 @@ extern "C" cuvsError_t cuvsHnswIndexParamsCreate(cuvsHnswIndexParams_t* params) .num_threads = 0, .M = 32, .metric = L2Expanded, - .ace_params = nullptr}; + .ace_params = nullptr, + .dataset_path = nullptr}; }); } @@ -404,3 +407,41 @@ extern "C" cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, } }); } + +extern "C" cuvsError_t cuvsHnswMaterializeParamsCreate(cuvsHnswMaterializeParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + *params = new cuvsHnswMaterializeParams{ + .dataset_path = nullptr, .max_host_memory_gb = 0, .num_threads = 0}; + }); +} + +extern "C" cuvsError_t cuvsHnswMaterializeParamsDestroy(cuvsHnswMaterializeParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} + +extern "C" cuvsError_t cuvsHnswMaterializeToHnswlib(cuvsResources_t res, + cuvsHnswMaterializeParams_t params, + const char* layered_artifact_path, + const char* output_path, + int dim, + cuvsDistanceType metric) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = reinterpret_cast(res); + auto cpp_params = cuvs::neighbors::hnsw::materialize_params(); + if (params->dataset_path != nullptr) { + cpp_params.dataset_path = std::string(params->dataset_path); + } + cpp_params.max_host_memory_gb = params->max_host_memory_gb; + cpp_params.num_threads = params->num_threads; + auto metric_type = static_cast(metric); + cuvs::neighbors::hnsw::materialize_to_hnswlib(*res_ptr, + cpp_params, + std::string(layered_artifact_path), + std::string(output_path), + dim, + metric_type); + }); +} diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 5800e05a96..bc5a58972a 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -1246,6 +1246,75 @@ void deserialize(raft::resources const& res, * @} */ +/** + * @defgroup hnsw_cpp_index_materialize Materialize a layered HNSW artifact into an hnswlib index + * @{ + */ + +/** + * @brief Parameters for materializing a layered HNSW artifact into an hnswlib index on disk. + */ +struct materialize_params { + /** Local dataset path holding the original-ID-ordered vectors used to build the artifact. + * + * Supported formats match layered deserialization: `.npy` and ANN benchmark `*.bin` files with a + * `[uint32 rows, uint32 cols]` header (`.fbin`, `.f16bin`, `.u8bin`, `.i8bin`). + */ + std::string dataset_path; + + /** Upper bound on host memory (in GiB) used for the base-topology reorder buffer. + * + * When `<= 0`, the whole base topology is reordered in a single in-memory pass (no temporary + * files). When set, the base topology is reordered through bucketed temporary files so that + * peak host memory stays close to this budget, at the cost of writing and re-reading the + * (small) base-topology section once. + */ + double max_host_memory_gb = 0; + + /** Number of host threads to use. When `0`, the maximum number of threads is used. */ + int num_threads = 0; +}; + +/** + * @brief Materialize a layered HNSW artifact into a standard hnswlib index file on disk. + * + * Materializes a `GPU_LAYERED_ON_DISK` artifact (graph topology only, stored in ACE order) plus a + * local dataset into a standard hnswlib index file, without ever holding the full materialized + * index in host memory. The materialization reorders the base topology from ACE order to + * original-id order and interleaves the vectors, emitting the output with sequential disk IO. The + * resulting file is compatible with the original hnswlib library (`loadIndex`) and can be read back + * through `cuvs::neighbors::hnsw::deserialize` with `hierarchy == HnswHierarchy::CPU`. + * + * The element data type (`float`, `half`, `uint8_t` or `int8_t`) is read from the artifact header, + * so a single entry point covers all supported dtypes. + * + * @param[in] res raft resources + * @param[in] params materialization parameters (dataset path, host-memory budget, threads) + * @param[in] layered_artifact_path path to the layered HNSW artifact + * @param[in] output_path path to the hnswlib index file to write + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric. Supported metrics ("L2Expanded", "InnerProduct") + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * hnsw::materialize_params materialize_params; + * materialize_params.dataset_path = "dataset.fbin"; + * hnsw::materialize_to_hnswlib( + * res, materialize_params, "layered_artifact.cuvs", "index.bin", dim, metric); + * @endcode + */ +void materialize_to_hnswlib(raft::resources const& res, + const materialize_params& params, + const std::string& layered_artifact_path, + const std::string& output_path, + int dim, + cuvs::distance::DistanceType metric); + +/** + * @} + */ + } // namespace hnsw } // namespace neighbors } // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index f02e9c2d7d..52d10a53e7 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -254,6 +254,19 @@ void write_large_file(const file_descriptor& fd, const size_t total_bytes, const uint64_t file_offset); +/** + * @brief Pre-size a file to `total_bytes` bytes. + * + * Prefers posix_fallocate (reserves blocks up-front, avoids later ENOSPC and fragmentation), but + * falls back to ftruncate on filesystems that do not support preallocation (tmpfs and some + * NFS/overlay mounts return EOPNOTSUPP/EINVAL/ENOSYS) so the operation still succeeds there. A + * `total_bytes` of 0 is a no-op. Throws on failure (uses the descriptor's path in the message). + * + * @param fd File descriptor to pre-size + * @param total_bytes Target file size in bytes + */ +void preallocate_file(const file_descriptor& fd, const size_t total_bytes); + /** * @brief Buffered output stream wrapper * diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 41fe241f98..f701ba962d 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -1032,11 +1033,7 @@ auto serialize_to_layered_hnsw_from_disk( const auto final_file_size = upper_links_offset + metadata.upper_links_bytes; cuvs::util::file_descriptor artifact_fd(artifact_file.string(), O_CREAT | O_RDWR | O_TRUNC, 0644); - const auto fallocate_result = posix_fallocate(artifact_fd.get(), 0, final_file_size); - RAFT_EXPECTS(fallocate_result == 0, - "Failed to pre-allocate layered HNSW artifact %s: %s", - artifact_file.string().c_str(), - std::strerror(fallocate_result)); + cuvs::util::preallocate_file(artifact_fd, final_file_size); cuvs::util::write_large_file(artifact_fd, &header, sizeof(header), 0); if (descriptors_bytes > 0) { cuvs::util::write_large_file( @@ -1176,6 +1173,45 @@ auto serialize_to_layered_hnsw_from_disk( return artifact_file.string(); } +// Build the standard hnswlib index header (matches HierarchicalNSW::saveIndex byte layout). Single +// source of the on-disk header encoding, shared by the streaming serialize path and the +// disk-to-disk materialize path so the two cannot drift. +inline auto make_hnswlib_native_header(size_t offset_level0, + size_t n_rows, + size_t size_data_per_element, + size_t label_offset, + size_t offset_data, + int maxlevel, + int enterpoint_node, + size_t maxM, + size_t maxM0, + size_t M, + double mult, + size_t ef_construction) -> std::vector +{ + std::vector buffer; + auto append = [&buffer](const auto& value) { + const auto* bytes = reinterpret_cast(&value); + buffer.insert(buffer.end(), bytes, bytes + sizeof(value)); + }; + const size_t max_elements = n_rows; + const size_t cur_element_count = n_rows; + append(offset_level0); + append(max_elements); + append(cur_element_count); + append(size_data_per_element); + append(label_offset); + append(offset_data); + append(maxlevel); + append(enterpoint_node); + append(maxM); + append(maxM0); + append(M); + append(mult); + append(ef_construction); + return buffer; +} + template void serialize_to_hnswlib_from_disk(raft::resources const& res, std::ostream& os_raw, @@ -1355,33 +1391,21 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, appr_algo->maxM0_, appr_algo->M_); - // offset_level_0 - os.write(reinterpret_cast(&appr_algo->offsetLevel0_), sizeof(std::size_t)); - // 8 max_element - override with n_rows - size_t num_elements = (size_t)n_rows; - os.write(reinterpret_cast(&num_elements), sizeof(std::size_t)); - // 16 curr_element_count - override with n_rows - os.write(reinterpret_cast(&num_elements), sizeof(std::size_t)); - // 24 size_data_per_element - os.write(reinterpret_cast(&appr_algo->size_data_per_element_), sizeof(std::size_t)); - // 32 label_offset - os.write(reinterpret_cast(&appr_algo->label_offset_), sizeof(std::size_t)); - // 40 offset_data - os.write(reinterpret_cast(&appr_algo->offsetData_), sizeof(std::size_t)); - // 48 maxlevel - os.write(reinterpret_cast(&appr_algo->maxlevel_), sizeof(int)); - // 52 enterpoint_node - os.write(reinterpret_cast(&appr_algo->enterpoint_node_), sizeof(int)); - // 56 maxM - os.write(reinterpret_cast(&appr_algo->maxM_), sizeof(std::size_t)); - // 64 maxM0 - os.write(reinterpret_cast(&appr_algo->maxM0_), sizeof(std::size_t)); - // 72 M - os.write(reinterpret_cast(&appr_algo->M_), sizeof(std::size_t)); - // 80 mult - os.write(reinterpret_cast(&appr_algo->mult_), sizeof(double)); - // 88 ef_construction - os.write(reinterpret_cast(&appr_algo->ef_construction_), sizeof(std::size_t)); + // Write the hnswlib index header (max_element / cur_element_count overridden with n_rows) using + // the shared encoder so the byte layout stays in lockstep with the materialize path. + const auto native_header = make_hnswlib_native_header(appr_algo->offsetLevel0_, + static_cast(n_rows), + appr_algo->size_data_per_element_, + appr_algo->label_offset_, + appr_algo->offsetData_, + appr_algo->maxlevel_, + appr_algo->enterpoint_node_, + appr_algo->maxM_, + appr_algo->maxM0_, + appr_algo->M_, + appr_algo->mult_, + appr_algo->ef_construction_); + os.write(native_header.data(), static_cast(native_header.size())); // host queries auto host_query_set = @@ -2060,17 +2084,32 @@ void serialize(raft::resources const& res, const std::string& filename, const in hnswlib_index->saveIndex(filename); } +// Parsed + fully validated view of a layered HNSW artifact: header metadata, the payload section +// offsets, and the per-row levels array. Both the in-memory deserialize path and the disk-to-disk +// materialize path consume this so the artifact format (offset arithmetic and the size-invariant +// checks) lives in a single place and cannot drift between them. +struct layered_artifact_view { + layered_hnsw_file_metadata metadata; + size_t artifact_size = 0; + size_t levels_offset = 0; + size_t base_nodes_offset = 0; + size_t base_links_offset = 0; + size_t upper_nodes_offset = 0; + size_t upper_links_offset = 0; + std::vector levels; // per-row top level, length n_rows +}; + +// Reads and validates the artifact header/descriptors/levels from an already-open artifact fd. +// Enforces every structural invariant (magic, version, dtype, metric, dim, layer count, section +// sizes, file length, and the levels max-level) BEFORE the caller allocates buffers sized from +// these fields, so a corrupt/truncated/crafted header cannot drive an out-of-bounds read or write. template -auto deserialize_layered_hnsw(raft::resources const& res, - const index_params& params, - const std::string& artifact_path, - int dim, - cuvs::distance::DistanceType metric) -> std::unique_ptr> +auto read_and_validate_layered_artifact(const cuvs::util::file_descriptor& artifact_fd, + const std::string& artifact_path, + int dim, + cuvs::distance::DistanceType metric) + -> layered_artifact_view { - common::nvtx::range fun_scope("hnsw::deserialize_layered"); - const auto total_start_time = std::chrono::steady_clock::now(); - const auto metadata_start_time = std::chrono::steady_clock::now(); - cuvs::util::file_descriptor artifact_fd(artifact_path, O_RDONLY); layered_hnsw_file_header header{}; cuvs::util::read_large_file(artifact_fd, &header, sizeof(header), 0); RAFT_EXPECTS(std::strncmp(header.magic, layered_hnsw_magic, sizeof(header.magic)) == 0, @@ -2088,23 +2127,25 @@ auto deserialize_layered_hnsw(raft::resources const& res, metric_name(static_cast(header.metric)), metric_name(metric)); - const auto artifact_size = static_cast(std::filesystem::file_size(artifact_path)); + layered_artifact_view view; + view.artifact_size = static_cast(std::filesystem::file_size(artifact_path)); const auto descriptors_offset = sizeof(layered_hnsw_file_header); const auto descriptors_bytes = static_cast(header.num_layers) * sizeof(layered_hnsw_layer_descriptor); - RAFT_EXPECTS(descriptors_offset + descriptors_bytes <= artifact_size, + RAFT_EXPECTS(descriptors_offset + descriptors_bytes <= view.artifact_size, "Layered HNSW layer descriptors are outside artifact: offset=%zu size=%zu " "artifact=%zu", descriptors_offset, descriptors_bytes, - artifact_size); + view.artifact_size); std::vector layer_descriptors(header.num_layers); if (descriptors_bytes > 0) { cuvs::util::read_large_file( artifact_fd, layer_descriptors.data(), descriptors_bytes, descriptors_offset); } - auto metadata = layered_hnsw_metadata_from_header(header); + auto& metadata = view.metadata; + metadata = layered_hnsw_metadata_from_header(header); metadata.layers.reserve(layer_descriptors.size()); for (const auto& descriptor : layer_descriptors) { metadata.layers.push_back({static_cast(descriptor.level), @@ -2113,7 +2154,6 @@ auto deserialize_layered_hnsw(raft::resources const& res, static_cast(descriptor.node_offset), static_cast(descriptor.link_offset)}); } - const auto metadata_elapsed_ms = elapsed_ms_since(metadata_start_time); RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); RAFT_EXPECTS(metadata.dim > 0, "Layered HNSW artifact must contain at least one dimension"); @@ -2126,21 +2166,70 @@ auto deserialize_layered_hnsw(raft::resources const& res, metadata.layers.size(), metadata.maxlevel); - RAFT_EXPECTS(!params.dataset_path.empty(), - "Layered HNSW deserialization requires index_params.dataset_path"); + // Section-size invariants: every buffer below is sized from n_rows / upper_nodes_count, so the + // declared section byte counts must match exactly or a later read/scatter would overrun. + RAFT_EXPECTS(metadata.levels_bytes == metadata.n_rows * sizeof(uint8_t), + "Layered HNSW levels section size mismatch"); + RAFT_EXPECTS(metadata.base_nodes_bytes == metadata.n_rows * sizeof(uint32_t), + "Layered HNSW base node section size mismatch"); + RAFT_EXPECTS(metadata.base_links_bytes == metadata.n_rows * metadata.base_link_row_bytes, + "Layered HNSW base links section size mismatch"); + RAFT_EXPECTS(metadata.upper_nodes_bytes == metadata.upper_nodes_count * sizeof(uint32_t), + "Layered HNSW upper node section size mismatch"); + RAFT_EXPECTS( + metadata.upper_links_bytes == metadata.upper_nodes_count * metadata.upper_link_row_bytes, + "Layered HNSW upper link section size mismatch"); + RAFT_EXPECTS(metadata.base_degree <= metadata.maxM0, + "Layered HNSW base degree (%zu) exceeds maxM0 (%zu)", + metadata.base_degree, + metadata.maxM0); const auto payload_offset = align_up(descriptors_offset + descriptors_bytes, layered_hnsw_alignment); - const auto levels_offset = payload_offset; - const auto base_nodes_offset = levels_offset + metadata.levels_bytes; - const auto base_links_offset = base_nodes_offset + metadata.base_nodes_bytes; - const auto upper_nodes_offset = base_links_offset + metadata.base_links_bytes; - const auto upper_links_offset = upper_nodes_offset + metadata.upper_nodes_bytes; - const auto expected_file_size = upper_links_offset + metadata.upper_links_bytes; - RAFT_EXPECTS(artifact_size >= expected_file_size, + view.levels_offset = payload_offset; + view.base_nodes_offset = view.levels_offset + metadata.levels_bytes; + view.base_links_offset = view.base_nodes_offset + metadata.base_nodes_bytes; + view.upper_nodes_offset = view.base_links_offset + metadata.base_links_bytes; + view.upper_links_offset = view.upper_nodes_offset + metadata.upper_nodes_bytes; + const auto expected_file_size = view.upper_links_offset + metadata.upper_links_bytes; + RAFT_EXPECTS(view.artifact_size >= expected_file_size, "Layered HNSW artifact is truncated: expected at least %zu bytes, got %zu", expected_file_size, - artifact_size); + view.artifact_size); + + view.levels.resize(metadata.n_rows); + cuvs::util::read_large_file( + artifact_fd, view.levels.data(), metadata.levels_bytes, view.levels_offset); + const auto max_level_in_levels = *std::max_element(view.levels.begin(), view.levels.end()); + RAFT_EXPECTS(static_cast(max_level_in_levels) == metadata.maxlevel, + "Layered HNSW levels max level (%d) does not match artifact maxlevel (%d)", + static_cast(max_level_in_levels), + metadata.maxlevel); + return view; +} + +template +auto deserialize_layered_hnsw(raft::resources const& res, + const index_params& params, + const std::string& artifact_path, + int dim, + cuvs::distance::DistanceType metric) -> std::unique_ptr> +{ + common::nvtx::range fun_scope("hnsw::deserialize_layered"); + const auto total_start_time = std::chrono::steady_clock::now(); + const auto metadata_start_time = std::chrono::steady_clock::now(); + cuvs::util::file_descriptor artifact_fd(artifact_path, O_RDONLY); + auto view = read_and_validate_layered_artifact(artifact_fd, artifact_path, dim, metric); + auto& metadata = view.metadata; + const auto artifact_size = view.artifact_size; + const auto base_nodes_offset = view.base_nodes_offset; + const auto base_links_offset = view.base_links_offset; + const auto upper_nodes_offset = view.upper_nodes_offset; + const auto upper_links_offset = view.upper_links_offset; + const auto metadata_elapsed_ms = elapsed_ms_since(metadata_start_time); + + RAFT_EXPECTS(!params.dataset_path.empty(), + "Layered HNSW deserialization requires index_params.dataset_path"); RAFT_LOG_INFO("Layered HNSW load: metadata read in %ld ms (rows=%zu dim=%zu artifact=%.2f GiB)", metadata_elapsed_ms, metadata.n_rows, @@ -2160,21 +2249,6 @@ auto deserialize_layered_hnsw(raft::resources const& res, dataset_file.shape.size() > 0 ? dataset_file.shape[0] : 0, dataset_file.shape.size() > 1 ? dataset_file.shape[1] : 0, params.dataset_path.c_str()); - RAFT_EXPECTS(metadata.levels_bytes == metadata.n_rows * sizeof(uint8_t), - "Layered HNSW levels section size mismatch"); - RAFT_EXPECTS(metadata.base_nodes_bytes == metadata.n_rows * sizeof(uint32_t), - "Layered HNSW base node section size mismatch"); - RAFT_EXPECTS(metadata.base_links_bytes == metadata.n_rows * metadata.base_link_row_bytes, - "Layered HNSW base links section size mismatch"); - RAFT_EXPECTS(metadata.upper_nodes_bytes == metadata.upper_nodes_count * sizeof(uint32_t), - "Layered HNSW upper node section size mismatch"); - RAFT_EXPECTS( - metadata.upper_links_bytes == metadata.upper_nodes_count * metadata.upper_link_row_bytes, - "Layered HNSW upper link section size mismatch"); - RAFT_EXPECTS(metadata.base_degree <= metadata.maxM0, - "Layered HNSW base degree (%zu) exceeds maxM0 (%zu)", - metadata.base_degree, - metadata.maxM0); RAFT_LOG_INFO("Layered HNSW load: dataset header validated in %ld ms (%s)", dataset_open_elapsed_ms, @@ -2202,19 +2276,8 @@ auto deserialize_layered_hnsw(raft::resources const& res, } }; - const auto levels_start_time = std::chrono::steady_clock::now(); - std::vector levels_u8(metadata.n_rows); - cuvs::util::read_large_file(artifact_fd, levels_u8.data(), metadata.levels_bytes, levels_offset); + auto& levels_u8 = view.levels; log_deserialize_progress(metadata.levels_bytes); - const auto max_level_in_levels = *std::max_element(levels_u8.begin(), levels_u8.end()); - RAFT_EXPECTS(static_cast(max_level_in_levels) == metadata.maxlevel, - "Layered HNSW levels max level (%d) does not match artifact maxlevel (%d)", - static_cast(max_level_in_levels), - metadata.maxlevel); - const auto levels_elapsed_ms = elapsed_ms_since(levels_start_time); - RAFT_LOG_INFO("Layered HNSW load: levels read in %ld ms (%.2f MiB)", - levels_elapsed_ms, - static_cast(metadata.levels_bytes) / (1024.0 * 1024.0)); const auto allocation_start_time = std::chrono::steady_clock::now(); auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); @@ -2427,6 +2490,668 @@ auto deserialize_layered_hnsw(raft::resources const& res, return hnsw_index; } +// Disk-to-disk materialization: layered HNSW artifact -> standard hnswlib index file. +// Constants/offsets shared by the materialization helpers. +struct hnswlib_materialize_layout { + size_t n_rows = 0; + size_t dim = 0; + size_t base_link_row_bytes = 0; + size_t upper_link_row_bytes = 0; + size_t size_data_per_element = 0; + size_t offset_data = 0; + size_t label_offset = 0; + size_t data_size = 0; + // Output file offsets. + size_t base_region_offset = 0; + size_t upper_region_offset = 0; + // Artifact payload offsets. + size_t base_nodes_offset = 0; + size_t base_links_offset = 0; + size_t upper_nodes_offset = 0; + size_t upper_links_offset = 0; +}; + +// Routes fixed-size, ID-keyed records into per-bucket temporary files and replays them per bucket. +// Used to reorder the (small) base/upper topology under a bounded host-memory budget while keeping +// all disk access sequential. +struct id_record_spiller { + std::filesystem::path dir; + size_t rows_per_bucket; + size_t num_buckets; + size_t record_bytes; + size_t buffer_cap; + std::vector fds; + std::vector offsets; + std::vector> buffers; + // One mutex per bucket so distinct buckets can be appended concurrently (add() is called from + // parallel scatter threads); same-bucket appends serialize on their own mutex. + std::vector bucket_mutexes; + + id_record_spiller(std::filesystem::path dir_, + size_t num_buckets_, + size_t rows_per_bucket_, + size_t record_bytes_, + size_t total_buffer_budget) + : dir(std::move(dir_)), + rows_per_bucket(rows_per_bucket_), + num_buckets(num_buckets_), + record_bytes(record_bytes_), + offsets(num_buckets_, 0), + buffers(num_buckets_), + bucket_mutexes(num_buckets_) + { + std::filesystem::create_directories(dir); + fds.reserve(num_buckets); + for (size_t b = 0; b < num_buckets; ++b) { + fds.emplace_back((dir / ("bucket_" + std::to_string(b) + ".tmp")).string(), + O_CREAT | O_RDWR | O_TRUNC, + 0644); + } + // Floor the per-bucket buffer so a tiny budget (many buckets) still batches many records per + // flush instead of degrading into one write syscall per record. This may exceed the nominal + // budget for pathologically small budgets, but keeps disk I/O coarse-grained and sequential. + constexpr size_t kMinRecordsPerFlush = 256; + const size_t even_share = total_buffer_budget / std::max(1, num_buckets); + buffer_cap = std::max(record_bytes * kMinRecordsPerFlush, even_share); + } + + id_record_spiller(const id_record_spiller&) = delete; + id_record_spiller& operator=(const id_record_spiller&) = delete; + id_record_spiller(id_record_spiller&&) = delete; + id_record_spiller& operator=(id_record_spiller&&) = delete; + + ~id_record_spiller() noexcept + { + fds.clear(); + std::error_code ec; + std::filesystem::remove_all(dir, ec); + } + + void flush(size_t b) + { + if (buffers[b].empty()) { return; } + cuvs::util::write_large_file(fds[b], buffers[b].data(), buffers[b].size(), offsets[b]); + offsets[b] += buffers[b].size(); + buffers[b].clear(); + } + + void add(size_t bucket, const void* record) + { + std::lock_guard guard(bucket_mutexes[bucket]); + auto& buf = buffers[bucket]; + if (buf.size() + record_bytes > buffer_cap) { flush(bucket); } + const auto* p = reinterpret_cast(record); + buf.insert(buf.end(), p, p + record_bytes); + } + + void finish_writes() + { + for (size_t b = 0; b < num_buckets; ++b) { + flush(b); + } + } + + template + void replay(size_t b, std::vector& chunk, F&& consume) + { + const size_t total = offsets[b]; + const size_t recs_per_chunk = std::max(1, (size_t{64} << 20) / record_bytes); + chunk.resize(recs_per_chunk * record_bytes); + size_t read_off = 0; + while (read_off < total) { + const size_t bytes = std::min(chunk.size(), total - read_off); + cuvs::util::read_large_file(fds[b], chunk.data(), bytes, read_off); + const size_t nrec = bytes / record_bytes; + for (size_t r = 0; r < nrec; ++r) { + consume(chunk.data() + r * record_bytes); + } + read_off += bytes; + } + } +}; + +// Streams the base topology (base_nodes + base_links) in artifact (ACE) order, invoking +// `cb(original_id, link_row_ptr)` for each row. When `parallel_threads > 1` the callback is invoked +// concurrently and must be thread-safe for distinct IDs. +template +void scatter_layered_base_links(const cuvs::util::file_descriptor& artifact_fd, + const hnswlib_materialize_layout& layout, + int parallel_threads, + F&& cb) +{ + const size_t n_rows = layout.n_rows; + const size_t row = layout.base_link_row_bytes; + const size_t batch = std::max(1, (size_t{64} << 20) / (sizeof(IdxT) + row)); + std::vector node_buf(batch); + std::vector link_buf(batch * row); + for (size_t s = 0; s < n_rows; s += batch) { + const size_t cur = std::min(batch, n_rows - s); + cuvs::util::read_large_file(artifact_fd, + node_buf.data(), + cur * sizeof(IdxT), + layout.base_nodes_offset + s * sizeof(IdxT)); + cuvs::util::read_large_file( + artifact_fd, link_buf.data(), cur * row, layout.base_links_offset + s * row); + bool invalid = false; + if (parallel_threads > 1) { +#pragma omp parallel for num_threads(parallel_threads) reduction(|| : invalid) + for (int64_t k = 0; k < static_cast(cur); ++k) { + const size_t id = static_cast(node_buf[k]); + if (id >= n_rows) { + invalid = true; + continue; + } + cb(id, link_buf.data() + static_cast(k) * row); + } + } else { + for (size_t k = 0; k < cur; ++k) { + const size_t id = static_cast(node_buf[k]); + if (id >= n_rows) { + invalid = true; + break; + } + cb(id, link_buf.data() + k * row); + } + } + RAFT_EXPECTS(!invalid, "Invalid base-layer node id in layered HNSW artifact"); + } +} + +// Streams the upper layers, invoking `cb(original_id, level, link_row_ptr)` for every promoted row. +template +void scatter_layered_upper_links(const cuvs::util::file_descriptor& artifact_fd, + const hnswlib_materialize_layout& layout, + const std::vector& layers, + const std::vector& levels_u8, + F&& cb) +{ + const size_t urow = layout.upper_link_row_bytes; + const size_t n_rows = layout.n_rows; + for (const auto& layer : layers) { + const size_t rc = layer.row_count; + if (rc == 0) { continue; } + const size_t batch = std::max(1, (size_t{64} << 20) / (sizeof(uint32_t) + urow)); + std::vector nodes(std::min(batch, rc)); + std::vector links(std::min(batch, rc) * urow); + for (size_t s = 0; s < rc; s += batch) { + const size_t cur = std::min(batch, rc - s); + cuvs::util::read_large_file( + artifact_fd, + nodes.data(), + cur * sizeof(uint32_t), + layout.upper_nodes_offset + (layer.node_offset + s) * sizeof(uint32_t)); + cuvs::util::read_large_file(artifact_fd, + links.data(), + cur * urow, + layout.upper_links_offset + (layer.link_offset + s) * urow); + for (size_t r = 0; r < cur; ++r) { + const size_t id = static_cast(nodes[r]); + RAFT_EXPECTS(id < n_rows, "Invalid upper-layer node id in layered HNSW artifact"); + RAFT_EXPECTS(layer.level <= static_cast(levels_u8[id]), + "Layered HNSW artifact references a node at an invalid upper level"); + cb(id, layer.level, links.data() + r * urow); + } + } + } +} + +// Emits the level-0 region `[link block | vector | label]` per ID, in increasing ID order, reading +// the dataset sequentially and writing the output sequentially. `row_ptr(id)` returns a pointer to +// the ID's level-0 link block in a caller-owned, bounded buffer. +template +void emit_hnswlib_base_records(const npy_file& dataset_file, + const cuvs::util::file_descriptor& output_fd, + const hnswlib_materialize_layout& layout, + int num_threads, + size_t id_begin, + size_t id_end, + std::vector& out_buffer, + DatasetBatch& dataset_batch, + RowPtr&& row_ptr) +{ + const size_t spe = layout.size_data_per_element; + const size_t row = layout.base_link_row_bytes; + const size_t data_size = layout.data_size; + const size_t out_batch_rows = static_cast(dataset_batch.extent(0)); + for (size_t s = id_begin; s < id_end; s += out_batch_rows) { + const size_t cur = std::min(out_batch_rows, id_end - s); + cuvs::util::read_large_file(dataset_file.fd, + dataset_batch.data_handle(), + cur * data_size, + dataset_file.header_size + s * data_size); +#pragma omp parallel for num_threads(num_threads) + for (int64_t k = 0; k < static_cast(cur); ++k) { + const size_t id = s + static_cast(k); + char* rec = out_buffer.data() + static_cast(k) * spe; + std::memcpy(rec, row_ptr(id), row); + std::memcpy(rec + layout.offset_data, + dataset_batch.data_handle() + static_cast(k) * layout.dim, + data_size); + const hnswlib::labeltype label = static_cast(id); + std::memcpy(rec + layout.label_offset, &label, sizeof(label)); + } + cuvs::util::write_large_file( + output_fd, out_buffer.data(), cur * spe, layout.base_region_offset + s * spe); + } +} + +// Phase 1 + 2: reorder the base topology to original-ID order and emit the level-0 region. +template +void materialize_hnswlib_base_region(const cuvs::util::file_descriptor& artifact_fd, + const npy_file& dataset_file, + const cuvs::util::file_descriptor& output_fd, + const layered_hnsw_file_metadata& metadata, + const hnswlib_materialize_layout& layout, + size_t budget_bytes, + int num_threads, + const std::filesystem::path& tmp_dir) +{ + const size_t n_rows = layout.n_rows; + const size_t row = layout.base_link_row_bytes; + + const size_t out_batch_rows = + std::max(1, (size_t{64} << 20) / layout.size_data_per_element); + auto dataset_batch = raft::make_host_matrix(static_cast(out_batch_rows), + static_cast(metadata.dim)); + std::vector out_buffer(out_batch_rows * layout.size_data_per_element); + + // Tracks that every original ID in [0, n_rows) is produced exactly once by the reorder. The + // section-size invariant (validated in Phase 0) guarantees exactly n_rows base-node records, so + // a malformed artifact with a missing/duplicate ID would otherwise leave a zero-initialized link + // row in the output and silently corrupt the graph; catch it instead of shipping a bad index. + std::vector seen(n_rows, 0); + auto verify_full_coverage = [&]() { + const auto covered = static_cast(std::count(seen.begin(), seen.end(), uint8_t{1})); + RAFT_EXPECTS(covered == n_rows, + "Layered HNSW base nodes cover %zu of %zu original ids (missing or duplicate ids)", + covered, + n_rows); + }; + + const auto base_start_time = std::chrono::steady_clock::now(); + // Single in-memory pass when the whole base topology fits the budget (no temporary files). + if (budget_bytes >= metadata.base_links_bytes) { + std::vector ordered_base(metadata.base_links_bytes); + scatter_layered_base_links( + artifact_fd, layout, num_threads, [&](size_t id, const char* link_row) { + std::memcpy(ordered_base.data() + id * row, link_row, row); + seen[id] = 1; + }); + verify_full_coverage(); + emit_hnswlib_base_records(dataset_file, + output_fd, + layout, + num_threads, + 0, + n_rows, + out_buffer, + dataset_batch, + [&](size_t id) { return ordered_base.data() + id * row; }); + RAFT_LOG_INFO("hnswlib materialize: base region written (single-pass) in %ld ms (%.2f GiB)", + elapsed_ms_since(base_start_time), + to_gib(layout.size_data_per_element * n_rows)); + return; + } + + // Bucketed reorder through temporary files for a hard memory budget. + const size_t record_bytes = sizeof(IdxT) + row; + size_t rows_per_bucket = std::max(1, budget_bytes / record_bytes); + rows_per_bucket = std::min(rows_per_bucket, n_rows); + const size_t num_buckets = (n_rows + rows_per_bucket - 1) / rows_per_bucket; + RAFT_LOG_INFO( + "hnswlib materialize: base region uses %zu buckets (rows/bucket=%zu, budget=%.2f GiB)", + num_buckets, + rows_per_bucket, + to_gib(budget_bytes)); + + id_record_spiller spiller( + tmp_dir / "base", num_buckets, rows_per_bucket, record_bytes, budget_bytes / 2); + // The spiller's per-bucket mutexes make add() safe under the parallel scatter; each thread stages + // the record in its own (thread_local) buffer before the routed append. + scatter_layered_base_links( + artifact_fd, layout, num_threads, [&](size_t id, const char* link_row) { + thread_local std::vector rec; + rec.resize(record_bytes); + const IdxT id32 = static_cast(id); + std::memcpy(rec.data(), &id32, sizeof(IdxT)); + std::memcpy(rec.data() + sizeof(IdxT), link_row, row); + spiller.add(id / rows_per_bucket, rec.data()); + seen[id] = 1; + }); + spiller.finish_writes(); + verify_full_coverage(); + + std::vector bucket_rows; + std::vector chunk; + for (size_t b = 0; b < num_buckets; ++b) { + const size_t id_begin = b * rows_per_bucket; + const size_t id_end = std::min(n_rows, id_begin + rows_per_bucket); + const size_t rib = id_end - id_begin; + bucket_rows.assign(rib * row, 0); + spiller.replay(b, chunk, [&](const char* r) { + IdxT id32; + std::memcpy(&id32, r, sizeof(IdxT)); + const size_t id = static_cast(id32); + std::memcpy(bucket_rows.data() + (id - id_begin) * row, r + sizeof(IdxT), row); + }); + emit_hnswlib_base_records( + dataset_file, + output_fd, + layout, + num_threads, + id_begin, + id_end, + out_buffer, + dataset_batch, + [&](size_t id) { return bucket_rows.data() + (id - id_begin) * row; }); + } + RAFT_LOG_INFO("hnswlib materialize: base region written (%zu buckets) in %ld ms (%.2f GiB)", + num_buckets, + elapsed_ms_since(base_start_time), + to_gib(layout.size_data_per_element * n_rows)); +} + +// Phase 3: transpose the upper layers into per-element link lists and emit the upper region. +inline void materialize_hnswlib_upper_region(const cuvs::util::file_descriptor& artifact_fd, + const cuvs::util::file_descriptor& output_fd, + const layered_hnsw_file_metadata& metadata, + const hnswlib_materialize_layout& layout, + const std::vector& levels_u8, + size_t budget_bytes, + const std::filesystem::path& tmp_dir) +{ + const size_t n_rows = layout.n_rows; + const size_t urow = layout.upper_link_row_bytes; + const auto upper_start_time = std::chrono::steady_clock::now(); + + // Sequential writer over the (variable-length) upper region. + size_t write_off = layout.upper_region_offset; + std::vector out_buffer; + out_buffer.reserve((size_t{64} << 20) + urow + sizeof(int)); + auto flush_out = [&]() { + if (out_buffer.empty()) { return; } + cuvs::util::write_large_file(output_fd, out_buffer.data(), out_buffer.size(), write_off); + write_off += out_buffer.size(); + out_buffer.clear(); + }; + auto append_element = [&](size_t level, const char* rows) { + const int link_list_size = level > 0 ? static_cast(level * urow) : 0; + const auto* p = reinterpret_cast(&link_list_size); + out_buffer.insert(out_buffer.end(), p, p + sizeof(int)); + if (level > 0) { out_buffer.insert(out_buffer.end(), rows, rows + level * urow); } + if (out_buffer.size() >= (size_t{64} << 20)) { flush_out(); } + }; + + const size_t fits_budget = metadata.upper_links_bytes + (n_rows + 1) * sizeof(size_t); + if (metadata.upper_links_bytes == 0 || fits_budget <= budget_bytes) { + // Pack all upper rows in original-ID order, then stream them out. + std::vector packed_start(n_rows + 1, 0); + for (size_t i = 0; i < n_rows; ++i) { + packed_start[i + 1] = packed_start[i] + static_cast(levels_u8[i]); + } + RAFT_EXPECTS(packed_start[n_rows] == metadata.upper_nodes_count, + "Layered HNSW upper rows (%zu) do not match upper_nodes_count (%zu)", + packed_start[n_rows], + metadata.upper_nodes_count); + std::vector packed(metadata.upper_links_bytes); + scatter_layered_upper_links(artifact_fd, + layout, + metadata.layers, + levels_u8, + [&](size_t id, size_t level, const char* link_row) { + const size_t dst_row = packed_start[id] + (level - 1); + std::memcpy(packed.data() + dst_row * urow, link_row, urow); + }); + size_t cursor = 0; + for (size_t id = 0; id < n_rows; ++id) { + const size_t level = static_cast(levels_u8[id]); + append_element(level, level > 0 ? packed.data() + cursor * urow : nullptr); + cursor += level; + } + flush_out(); + RAFT_LOG_INFO("hnswlib materialize: upper region written (single-pass) in %ld ms (%.2f GiB)", + elapsed_ms_since(upper_start_time), + to_gib(metadata.upper_links_bytes)); + return; + } + + // Bucketed upper transpose for a hard memory budget. + const size_t budget_half = std::max(1, budget_bytes / 2); + size_t num_buckets = + std::max(1, (metadata.upper_links_bytes + budget_half - 1) / budget_half); + size_t rows_per_bucket = std::max(1, (n_rows + num_buckets - 1) / num_buckets); + num_buckets = (n_rows + rows_per_bucket - 1) / rows_per_bucket; + RAFT_LOG_INFO("hnswlib materialize: upper region uses %zu buckets (rows/bucket=%zu)", + num_buckets, + rows_per_bucket); + + const size_t record_bytes = 2 * sizeof(uint32_t) + urow; // [id][level][row] + id_record_spiller spiller( + tmp_dir / "upper", num_buckets, rows_per_bucket, record_bytes, budget_half); + std::vector rec(record_bytes); + scatter_layered_upper_links(artifact_fd, + layout, + metadata.layers, + levels_u8, + [&](size_t id, size_t level, const char* link_row) { + const uint32_t id32 = static_cast(id); + const uint32_t lv32 = static_cast(level); + std::memcpy(rec.data(), &id32, sizeof(uint32_t)); + std::memcpy(rec.data() + sizeof(uint32_t), &lv32, sizeof(uint32_t)); + std::memcpy(rec.data() + 2 * sizeof(uint32_t), link_row, urow); + spiller.add(id / rows_per_bucket, rec.data()); + }); + spiller.finish_writes(); + + std::vector chunk; + for (size_t b = 0; b < num_buckets; ++b) { + const size_t id_begin = b * rows_per_bucket; + const size_t id_end = std::min(n_rows, id_begin + rows_per_bucket); + const size_t rib = id_end - id_begin; + std::vector local_start(rib + 1, 0); + for (size_t i = 0; i < rib; ++i) { + local_start[i + 1] = local_start[i] + static_cast(levels_u8[id_begin + i]); + } + std::vector packed(local_start[rib] * urow); + spiller.replay(b, chunk, [&](const char* r) { + uint32_t id32 = 0; + uint32_t lv32 = 0; + std::memcpy(&id32, r, sizeof(uint32_t)); + std::memcpy(&lv32, r + sizeof(uint32_t), sizeof(uint32_t)); + const size_t local = static_cast(id32) - id_begin; + const size_t dst_row = local_start[local] + (static_cast(lv32) - 1); + std::memcpy(packed.data() + dst_row * urow, r + 2 * sizeof(uint32_t), urow); + }); + size_t cursor = 0; + for (size_t i = 0; i < rib; ++i) { + const size_t level = static_cast(levels_u8[id_begin + i]); + append_element(level, level > 0 ? packed.data() + cursor * urow : nullptr); + cursor += level; + } + } + flush_out(); + RAFT_LOG_INFO("hnswlib materialize: upper region written (%zu buckets) in %ld ms (%.2f GiB)", + num_buckets, + elapsed_ms_since(upper_start_time), + to_gib(metadata.upper_links_bytes)); +} + +// Materialize a layered HNSW artifact + dataset into a standard hnswlib index file on disk, using +// bounded host memory and sequential disk I/O. +template +void materialize_layered_to_hnswlib_on_disk(raft::resources const& res, + const cuvs::util::file_descriptor& artifact_fd, + const cuvs::neighbors::hnsw::materialize_params& params, + const std::string& layered_artifact_path, + const std::string& output_path, + int dim, + cuvs::distance::DistanceType metric) +{ + static_assert(std::is_same_v, "Layered HNSW artifacts store ids as uint32_t"); + common::nvtx::range fun_scope("hnsw::materialize_to_hnswlib"); + const auto total_start_time = std::chrono::steady_clock::now(); + + // ---- Phase 0: read and validate the artifact header, descriptors and levels ---- + // Shared with the in-memory deserialize path so the format/offset arithmetic and every + // size-invariant check live in one place; this also bounds the buffer allocations below. + auto view = + read_and_validate_layered_artifact(artifact_fd, layered_artifact_path, dim, metric); + auto& metadata = view.metadata; + auto& levels_u8 = view.levels; + const auto base_nodes_offset = view.base_nodes_offset; + const auto base_links_offset = view.base_links_offset; + const auto upper_nodes_offset = view.upper_nodes_offset; + const auto upper_links_offset = view.upper_links_offset; + RAFT_EXPECTS(!params.dataset_path.empty(), + "Layered HNSW materialization requires materialize_params.dataset_path"); + + const size_t n_rows = metadata.n_rows; + + // Validate the dataset header. + auto dataset_file = open_layered_dataset_file(params.dataset_path); + RAFT_EXPECTS(dataset_file.shape.size() == 2 && dataset_file.shape[0] == n_rows && + dataset_file.shape[1] == metadata.dim, + "Layered HNSW dataset shape mismatch: artifact rows=%zu dim=%zu, dataset path=%s", + n_rows, + metadata.dim, + params.dataset_path.c_str()); + + // Retrieve the hnswlib layout constants from a dummy (single-element) index. + auto hnsw_index = std::make_unique>(dim, metric, HnswHierarchy::CPU); + auto appr_algo = std::make_unique::type>>( + hnsw_index->get_space(), 1, metadata.M, metadata.ef_construction); + const size_t data_size = static_cast(dim) * sizeof(T); + RAFT_EXPECTS(appr_algo->size_links_level0_ == metadata.base_link_row_bytes, + "Layered HNSW base link row size mismatch"); + RAFT_EXPECTS(appr_algo->size_links_per_element_ == metadata.upper_link_row_bytes, + "Layered HNSW upper link row size mismatch"); + RAFT_EXPECTS(appr_algo->maxM0_ == metadata.maxM0 && appr_algo->maxM_ == metadata.maxM, + "Layered HNSW M parameter mismatch"); + RAFT_EXPECTS(appr_algo->data_size_ == data_size, "Layered HNSW data size mismatch"); + RAFT_EXPECTS(appr_algo->offsetData_ == appr_algo->size_links_level0_, + "Unexpected hnswlib data offset"); + RAFT_EXPECTS(appr_algo->label_offset_ == appr_algo->size_links_level0_ + data_size, + "Unexpected hnswlib label offset"); + + // ---- Output layout ---- + const auto native_header = make_hnswlib_native_header(appr_algo->offsetLevel0_, + n_rows, + appr_algo->size_data_per_element_, + appr_algo->label_offset_, + appr_algo->offsetData_, + metadata.maxlevel, + metadata.enterpoint_node, + metadata.maxM, + metadata.maxM0, + metadata.M, + metadata.mult, + metadata.ef_construction); + hnswlib_materialize_layout layout; + layout.n_rows = n_rows; + layout.dim = metadata.dim; + layout.base_link_row_bytes = metadata.base_link_row_bytes; + layout.upper_link_row_bytes = metadata.upper_link_row_bytes; + layout.size_data_per_element = appr_algo->size_data_per_element_; + layout.offset_data = appr_algo->offsetData_; + layout.label_offset = appr_algo->label_offset_; + layout.data_size = data_size; + layout.base_region_offset = native_header.size(); + layout.base_nodes_offset = base_nodes_offset; + layout.base_links_offset = base_links_offset; + layout.upper_nodes_offset = upper_nodes_offset; + layout.upper_links_offset = upper_links_offset; + + const size_t base_region_bytes = n_rows * layout.size_data_per_element; + layout.upper_region_offset = layout.base_region_offset + base_region_bytes; + + size_t upper_region_bytes = 0; + for (size_t i = 0; i < n_rows; ++i) { + upper_region_bytes += sizeof(int); + const size_t level = static_cast(levels_u8[i]); + if (level > 0) { upper_region_bytes += level * layout.upper_link_row_bytes; } + } + const size_t final_size = layout.upper_region_offset + upper_region_bytes; + + cuvs::util::file_descriptor output_fd(output_path, O_CREAT | O_RDWR | O_TRUNC, 0644); + cuvs::util::preallocate_file(output_fd, final_size); + cuvs::util::write_large_file(output_fd, native_header.data(), native_header.size(), 0); + RAFT_LOG_INFO( + "hnswlib materialize: writing index (rows=%zu dim=%zu, output=%.2f GiB, dataset=%s)", + n_rows, + metadata.dim, + to_gib(final_size), + params.dataset_path.c_str()); + + auto num_threads = + params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; + // max_host_memory_gb <= 0 => no host-memory cap: both regions take their single in-memory pass + // and create no temporary files. Otherwise reserve headroom for the fixed-size streaming I/O + // buffers (dataset batch, output buffer, per-call scatter read buffers, each ~64 MiB) so peak + // host memory stays close to the requested budget rather than overshooting it by the buffer + // overhead. + size_t budget_bytes = std::numeric_limits::max(); + if (params.max_host_memory_gb > 0.0) { + const size_t requested = + std::max(1, static_cast(params.max_host_memory_gb * (size_t{1} << 30))); + constexpr size_t kStreamingBufferReserve = size_t{4} * (size_t{64} << 20); // ~256 MiB + budget_bytes = requested > kStreamingBufferReserve ? requested - kStreamingBufferReserve + : std::max(1, requested / 2); + } + + std::filesystem::path tmp_dir = output_path; + tmp_dir += ".materialize_tmp"; + + materialize_hnswlib_base_region( + artifact_fd, dataset_file, output_fd, metadata, layout, budget_bytes, num_threads, tmp_dir); + materialize_hnswlib_upper_region( + artifact_fd, output_fd, metadata, layout, levels_u8, budget_bytes, tmp_dir); + + std::error_code ec; + std::filesystem::remove_all(tmp_dir, ec); + + RAFT_LOG_INFO("hnswlib materialize: completed in %ld ms (output %.2f GiB)", + elapsed_ms_since(total_start_time), + to_gib(final_size)); +} + +// Reads the dtype tag from the artifact header and dispatches to the typed implementation. +inline void materialize_layered_to_hnswlib_on_disk_dispatch( + raft::resources const& res, + const cuvs::neighbors::hnsw::materialize_params& params, + const std::string& layered_artifact_path, + const std::string& output_path, + int dim, + cuvs::distance::DistanceType metric) +{ + // Open the artifact once and read just the dtype tag to pick T; the typed worker reuses this fd + // (no second open) and performs full header validation via read_and_validate_layered_artifact. + cuvs::util::file_descriptor artifact_fd(layered_artifact_path, O_RDONLY); + layered_hnsw_file_header header{}; + cuvs::util::read_large_file(artifact_fd, &header, sizeof(header), 0); + RAFT_EXPECTS(std::strncmp(header.magic, layered_hnsw_magic, sizeof(header.magic)) == 0, + "Invalid layered HNSW artifact magic: %s", + layered_artifact_path.c_str()); + switch (static_cast(header.dtype)) { + case layered_hnsw_dtype::float32: + materialize_layered_to_hnswlib_on_disk( + res, artifact_fd, params, layered_artifact_path, output_path, dim, metric); + break; + case layered_hnsw_dtype::float16: + materialize_layered_to_hnswlib_on_disk( + res, artifact_fd, params, layered_artifact_path, output_path, dim, metric); + break; + case layered_hnsw_dtype::uint8: + materialize_layered_to_hnswlib_on_disk( + res, artifact_fd, params, layered_artifact_path, output_path, dim, metric); + break; + case layered_hnsw_dtype::int8: + materialize_layered_to_hnswlib_on_disk( + res, artifact_fd, params, layered_artifact_path, output_path, dim, metric); + break; + default: RAFT_FAIL("Unsupported layered HNSW artifact dtype %u", header.dtype); + } +} + template void deserialize(raft::resources const& res, const index_params& params, diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp index 54e9dcf12a..a3d6ca3ab6 100644 --- a/cpp/src/neighbors/hnsw.cpp +++ b/cpp/src/neighbors/hnsw.cpp @@ -115,4 +115,17 @@ CUVS_INST_HNSW_SERIALIZE(int8_t); #undef CUVS_INST_HNSW_SERIALIZE +// The element data type is read from the artifact header; the dispatcher selects the typed +// implementation, instantiating the float/half/uint8_t/int8_t materialize paths in this TU. +void materialize_to_hnswlib(raft::resources const& res, + const materialize_params& params, + const std::string& layered_artifact_path, + const std::string& output_path, + int dim, + cuvs::distance::DistanceType metric) +{ + detail::materialize_layered_to_hnswlib_on_disk_dispatch( + res, params, layered_artifact_path, output_path, dim, metric); +} + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index d924527e72..dac930cad2 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -10,10 +10,29 @@ #include #include +#include #include namespace cuvs::util { +void preallocate_file(const file_descriptor& fd, const size_t total_bytes) +{ + if (total_bytes == 0) { return; } + RAFT_EXPECTS(fd.is_valid(), "File descriptor must be valid"); + const int rc = posix_fallocate(fd.get(), 0, static_cast(total_bytes)); + if (rc == 0) { return; } + // Some filesystems (tmpfs, certain NFS/overlay mounts) do not support preallocation; fall back + // to ftruncate so a valid output location is still usable. + if (rc == EOPNOTSUPP || rc == EINVAL || rc == ENOSYS) { + RAFT_EXPECTS(ftruncate(fd.get(), static_cast(total_bytes)) == 0, + "Failed to pre-size file %s via ftruncate: %s", + fd.get_path().c_str(), + strerror(errno)); + return; + } + RAFT_FAIL("Failed to pre-allocate file %s: %s", fd.get_path().c_str(), strerror(rc)); +} + void read_large_file(const file_descriptor& fd, void* dest_ptr, const size_t total_bytes, diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 46caef7ae1..f0eeaa51ad 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -16,7 +16,9 @@ #include #include #include +#include #include +#include namespace cuvs::neighbors::hnsw { @@ -498,6 +500,206 @@ class AnnHnswAceTest : public ::testing::TestWithParam { std::exception); } + void testHnswAceLayeredMaterializeToHnswlib() + { + size_t queries_size = ps.n_queries * ps.k; + std::vector indexes_naive(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indexes_naive_dev(queries_size, stream_); + cuvs::neighbors::naive_knn(handle_, + distances_naive_dev.data(), + indexes_naive_dev.data(), + search_queries.data(), + database_dev.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::update_host(indexes_naive.data(), indexes_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + std::string temp_dir = std::string("/tmp/cuvs_hnsw_ace_materialize_test_") + + std::to_string(std::time(nullptr)) + "_" + + std::to_string(reinterpret_cast(this)); + std::filesystem::create_directories(temp_dir); + struct temp_dir_cleanup { + std::string path; + ~temp_dir_cleanup() + { + std::error_code ec; + std::filesystem::remove_all(path, ec); + } + } cleanup{temp_dir}; + + auto database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); + auto queries_host = raft::make_host_matrix(ps.n_queries, ps.dim); + raft::copy(queries_host.data_handle(), search_queries.data(), ps.n_queries * ps.dim, stream_); + raft::resource::sync_stream(handle_); + + const auto dataset_file = (std::filesystem::path(temp_dir) / "dataset.npy").string(); + auto [dataset_fd, dataset_header_size] = cuvs::util::create_numpy_file( + dataset_file, {static_cast(ps.n_rows), static_cast(ps.dim)}); + cuvs::util::write_large_file(dataset_fd, + database_host.data_handle(), + static_cast(ps.n_rows) * ps.dim * sizeof(DataT), + dataset_header_size); + + hnsw::index_params hnsw_params; + hnsw_params.metric = ps.metric; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU_LAYERED_ON_DISK; + hnsw_params.M = 32; + hnsw_params.ef_construction = ps.ef_construction; + hnsw_params.dataset_path = dataset_file; + + auto ace_params = graph_build_params::ace_params(); + ace_params.npartitions = ps.npartitions; + ace_params.build_dir = temp_dir; + ace_params.use_disk = true; + ace_params.max_host_memory_gb = ps.max_host_memory_gb; + ace_params.max_gpu_memory_gb = ps.max_gpu_memory_gb; + hnsw_params.graph_build_params = ace_params; + + auto hnsw_index = + hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); + ASSERT_NE(hnsw_index, nullptr); + const auto artifact_path = hnsw_index->file_path(); + ASSERT_FALSE(artifact_path.empty()); + + hnsw::search_params search_params; + search_params.ef = std::max(ps.ef_construction, ps.k * 2); + search_params.num_threads = 1; + + auto indexes_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + auto distances_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + + // Reference: load the layered artifact in RAM and search. + std::vector indexes_layered(queries_size); + std::vector distances_layered(queries_size); + { + hnsw::index* layered_index = nullptr; + hnsw::deserialize(handle_, hnsw_params, artifact_path, ps.dim, ps.metric, &layered_index); + ASSERT_NE(layered_index, nullptr); + std::unique_ptr> layered_guard(layered_index); + hnsw::search(handle_, + search_params, + *layered_guard, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + for (size_t i = 0; i < queries_size; i++) { + indexes_layered[i] = static_cast(indexes_hnsw_host.data_handle()[i]); + distances_layered[i] = distances_hnsw_host.data_handle()[i]; + } + } + + // Materialize to a standard hnswlib index file twice: single in-memory pass and bucketed. + const auto out_single = (std::filesystem::path(temp_dir) / "materialized_single.bin").string(); + const auto out_bucketed = + (std::filesystem::path(temp_dir) / "materialized_bucketed.bin").string(); + + hnsw::materialize_params materialize_single; + materialize_single.dataset_path = dataset_file; + materialize_single.max_host_memory_gb = 0; // single in-memory reorder pass + hnsw::materialize_to_hnswlib( + handle_, materialize_single, artifact_path, out_single, ps.dim, ps.metric); + + hnsw::materialize_params materialize_bucketed; + materialize_bucketed.dataset_path = dataset_file; + materialize_bucketed.max_host_memory_gb = 0.00003; // tiny budget forces bucketed base + upper + hnsw::materialize_to_hnswlib( + handle_, materialize_bucketed, artifact_path, out_bucketed, ps.dim, ps.metric); + + ASSERT_TRUE(std::filesystem::is_regular_file(out_single)); + ASSERT_TRUE(std::filesystem::is_regular_file(out_bucketed)); + + // Determinism: single-pass and bucketed outputs must be byte-identical. + { + const auto read_all = [](const std::string& path) { + std::ifstream stream(path, std::ios::binary); + return std::vector((std::istreambuf_iterator(stream)), + std::istreambuf_iterator()); + }; + const auto bytes_single = read_all(out_single); + const auto bytes_bucketed = read_all(out_bucketed); + ASSERT_FALSE(bytes_single.empty()); + EXPECT_EQ(bytes_single.size(), bytes_bucketed.size()); + EXPECT_TRUE(bytes_single == bytes_bucketed) + << "Single-pass and bucketed materialized outputs differ"; + } + + // Load the materialized file as a standard (CPU) hnswlib index and search. + hnsw::index_params cpu_params; + cpu_params.metric = ps.metric; + cpu_params.hierarchy = hnsw::HnswHierarchy::CPU; + hnsw::index* cpu_index = nullptr; + hnsw::deserialize(handle_, cpu_params, out_single, ps.dim, ps.metric, &cpu_index); + ASSERT_NE(cpu_index, nullptr); + std::unique_ptr> cpu_guard(cpu_index); + + hnsw::search(handle_, + search_params, + *cpu_guard, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + + std::vector indexes_cpu(queries_size); + std::vector distances_cpu(queries_size); + for (size_t i = 0; i < queries_size; i++) { + indexes_cpu[i] = static_cast(indexes_hnsw_host.data_handle()[i]); + distances_cpu[i] = distances_hnsw_host.data_handle()[i]; + } + + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_naive, + indexes_cpu, + distances_naive, + distances_cpu, + ps.n_queries, + ps.k, + 0.003, + ps.min_recall)) + << "Materialized hnswlib index failed recall check vs. ground truth"; + + // The materialized CPU index represents the same graph and vectors as the layered artifact, + // so its search results must match the in-memory layered path. + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_layered, + indexes_cpu, + distances_layered, + distances_cpu, + ps.n_queries, + ps.k, + 0.003, + 0.99)) + << "Materialized hnswlib index disagrees with the in-memory layered path"; + + hnsw::materialize_params bad_params; + bad_params.dataset_path = dataset_file; + const auto err_out = (std::filesystem::path(temp_dir) / "err.bin").string(); + EXPECT_THROW(hnsw::materialize_to_hnswlib( + handle_, bad_params, artifact_path, err_out, ps.dim + 1, ps.metric), + std::exception); + + const auto wrong_metric = ps.metric == cuvs::distance::DistanceType::L2Expanded + ? cuvs::distance::DistanceType::InnerProduct + : cuvs::distance::DistanceType::L2Expanded; + EXPECT_THROW(hnsw::materialize_to_hnswlib( + handle_, bad_params, artifact_path, err_out, ps.dim, wrong_metric), + std::exception); + + hnsw::materialize_params missing_dataset; + missing_dataset.dataset_path.clear(); + EXPECT_THROW(hnsw::materialize_to_hnswlib( + handle_, missing_dataset, artifact_path, err_out, ps.dim, ps.metric), + std::exception); + } + void SetUp() override { database_dev.resize(((size_t)ps.n_rows) * ps.dim, stream_); diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu index 97a6255300..33a272a675 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -33,4 +33,14 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, AnnHnswAceLayeredTest_float, ::testing::ValuesIn(hnsw_ace_layered_inputs)); +typedef AnnHnswAceTest AnnHnswAceMaterializeTest_float; +TEST_P(AnnHnswAceMaterializeTest_float, AnnHnswAceLayeredMaterializeToHnswlib) +{ + this->testHnswAceLayeredMaterializeToHnswlib(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMaterializeTest, + AnnHnswAceMaterializeTest_float, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu index c77391ba28..acde405a3b 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -32,4 +32,14 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, AnnHnswAceLayeredTest_half, ::testing::ValuesIn(hnsw_ace_layered_inputs)); +typedef AnnHnswAceTest AnnHnswAceMaterializeTest_half; +TEST_P(AnnHnswAceMaterializeTest_half, AnnHnswAceLayeredMaterializeToHnswlib) +{ + this->testHnswAceLayeredMaterializeToHnswlib(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMaterializeTest, + AnnHnswAceMaterializeTest_half, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu index 4ee6ecc5e5..a6fe986898 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -34,4 +34,14 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, AnnHnswAceLayeredTest_int8_t, ::testing::ValuesIn(hnsw_ace_layered_inputs)); +typedef AnnHnswAceTest AnnHnswAceMaterializeTest_int8_t; +TEST_P(AnnHnswAceMaterializeTest_int8_t, AnnHnswAceLayeredMaterializeToHnswlib) +{ + this->testHnswAceLayeredMaterializeToHnswlib(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMaterializeTest, + AnnHnswAceMaterializeTest_int8_t, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu index 1ef791df8e..13c553f260 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -34,4 +34,14 @@ INSTANTIATE_TEST_CASE_P(AnnHnswAceLayeredTest, AnnHnswAceLayeredTest_uint8_t, ::testing::ValuesIn(hnsw_ace_layered_inputs)); +typedef AnnHnswAceTest AnnHnswAceMaterializeTest_uint8_t; +TEST_P(AnnHnswAceMaterializeTest_uint8_t, AnnHnswAceLayeredMaterializeToHnswlib) +{ + this->testHnswAceLayeredMaterializeToHnswlib(); +} + +INSTANTIATE_TEST_CASE_P(AnnHnswAceMaterializeTest, + AnnHnswAceMaterializeTest_uint8_t, + ::testing::ValuesIn(hnsw_ace_layered_inputs)); + } // namespace cuvs::neighbors::hnsw diff --git a/examples/cpp/src/hnsw_ace_layered_example.cu b/examples/cpp/src/hnsw_ace_layered_example.cu index f11551df20..ced22ff3fe 100644 --- a/examples/cpp/src/hnsw_ace_layered_example.cu +++ b/examples/cpp/src/hnsw_ace_layered_example.cu @@ -11,12 +11,15 @@ // disk. It emits one topology-only artifact, hnsw_index.cuvs. The dataset remains separate and does // not need to be transferred to the search server, which typically has the dataset locally. // -// This example demonstrates how to build a layered HNSW index with ACE: +// This example demonstrates how to build a layered HNSW index with ACE and turn it into a standard +// hnswlib index for in-memory search: // // 1. Optionally quantize the dataset and queries to int8. // 2. Build a single-file layered HNSW artifact with ACE using hnsw::build. -// 3. Deserialize the layered HNSW artifact using hnsw::deserialize. -// 4. Search the in-memory HNSW index. +// 3. Materialize the layered artifact into a standard hnswlib index file on disk using +// hnsw::materialize_to_hnswlib (disk-to-disk, never holding the full index in host memory). +// 4. Read the materialized hnswlib index into memory using hnsw::deserialize (hierarchy = CPU). +// 5. Search the in-memory HNSW index. // // Layered-on-disk layout: // @@ -27,14 +30,14 @@ // upper nodes + upper links: hnswlib-ready upper-layer topology // // The transferred index artifact is topology-only. The dataset is loaded locally during -// deserialization from hnsw::index_params::dataset_path. The loader supports .npy and ANN benchmark -// *.bin datasets; this example writes a local dataset .npy only to make the demo self-contained. -// -// Layer 0 node IDs and neighbor IDs are original dataset row IDs. Upper layers are generated with -// the same level/order/KNN logic as serialize_to_hnswlib_from_disk, then stored as hnswlib-ready -// link rows so deserialization does no graph remapping or link padding on the search node. +// materialization from hnsw::materialize_params::dataset_path. The loader supports .npy and ANN +// benchmark *.bin datasets; this example writes a local dataset .npy only to make the demo +// self-contained. The materialized hnswlib index file is self-contained (it embeds the vectors), +// so reading it back needs no dataset path. #include +#include +#include #include #include #include @@ -46,6 +49,7 @@ #include #include #include +#include #include #include @@ -64,6 +68,18 @@ namespace { constexpr const char* kBuildDir = "/tmp/hnsw_ace_layered"; +// Reports the wall-clock time of a callable in milliseconds. +template +double time_ms(F&& fn) +{ + const auto start = std::chrono::steady_clock::now(); + fn(); + return std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); +} + +double to_gib(double bytes) { return bytes / (1024.0 * 1024.0 * 1024.0); } + template std::string write_local_dataset(raft::host_matrix_view dataset, const std::string& path) @@ -130,29 +146,72 @@ auto hnsw_build(raft::device_resources const& dev_resources, { using namespace cuvs::neighbors; - auto hnsw_index = hnsw::build(dev_resources, hnsw_params, dataset); + std::unique_ptr> hnsw_index; + const auto build_ms = + time_ms([&]() { hnsw_index = hnsw::build(dev_resources, hnsw_params, dataset); }); const auto artifact_path = hnsw_index->file_path(); if (artifact_path.empty()) { throw std::runtime_error("Expected layered HNSW build to return an artifact path."); } - std::cout << " hnsw_build: layered artifact written to " << artifact_path << std::endl; + const auto artifact_bytes = static_cast(std::filesystem::file_size(artifact_path)); + std::cout << " hnsw_build: layered artifact written to " << artifact_path << "\n" + << " hnsw_build: build wall time " << build_ms << " ms, artifact " + << to_gib(artifact_bytes) << " GiB" << std::endl; return artifact_path; } +// Materialize the layered artifact into a standard hnswlib index file on disk and time the +// disk-to-disk materialization. Returns the path to the native hnswlib index file. template -auto hnsw_deserialize(raft::device_resources const& dev_resources, +auto hnsw_materialize(raft::device_resources const& dev_resources, const cuvs::neighbors::hnsw::index_params& hnsw_params, const std::string& artifact_path, + const std::string& dataset_path, + int64_t dim, + const std::string& output_path) -> std::string +{ + using namespace cuvs::neighbors; + + hnsw::materialize_params materialize_params; + materialize_params.dataset_path = dataset_path; + materialize_params.max_host_memory_gb = 0; // 0 => single in-memory reorder pass + materialize_params.num_threads = 0; // 0 => max threads + + const auto materialize_ms = time_ms([&]() { + hnsw::materialize_to_hnswlib(dev_resources, + materialize_params, + artifact_path, + output_path, + static_cast(dim), + hnsw_params.metric); + }); + + const auto native_bytes = static_cast(std::filesystem::file_size(output_path)); + std::cout << " hnsw_materialize: native hnswlib index written to " << output_path << "\n" + << " hnsw_materialize: wall time " << materialize_ms << " ms, output " + << to_gib(native_bytes) << " GiB" << std::endl; + return output_path; +} + +// Read the materialized hnswlib index into memory for search. The materialized file is a standard +// hnswlib index, so it is loaded with hierarchy == CPU and needs no dataset path (the file already +// embeds the vectors). +template +auto hnsw_load_native(raft::device_resources const& dev_resources, + const std::string& native_index_path, + cuvs::distance::DistanceType metric, int64_t dim) -> std::unique_ptr> { using namespace cuvs::neighbors; - hnsw::index* deserialized_index = nullptr; - // Set params.dataset_path to the local dataset path to load the dataset from disk. - // hnsw_params.dataset_path = "/tmp/dataset.npy"; + hnsw::index_params load_params; + load_params.hierarchy = hnsw::HnswHierarchy::CPU; + load_params.metric = metric; + + hnsw::index* loaded_index = nullptr; hnsw::deserialize( - dev_resources, hnsw_params, artifact_path, dim, hnsw_params.metric, &deserialized_index); - return std::unique_ptr>(deserialized_index); + dev_resources, load_params, native_index_path, static_cast(dim), metric, &loaded_index); + return std::unique_ptr>(loaded_index); } template @@ -206,6 +265,9 @@ int main() { raft::device_resources dev_resources; + // Surface the per-phase build/materialize timing logs (RAFT_LOG_INFO). + raft::default_logger().set_level(rapids_logger::level_enum::info); + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), 1024 * 1024 * 1024ull); rmm::mr::set_current_device_resource(pool_mr); @@ -251,26 +313,40 @@ int main() auto dataset_path = write_local_dataset(dataset_i8_view, std::string{kBuildDir} + "/dataset.npy"); auto hnsw_params = make_hnsw_ace_params(kBuildDir, dataset_path); + const std::string native_index_path = std::string{kBuildDir} + "/hnsw_native.bin"; + std::cout << "[stage 2] Build layered HNSW index with ACE" << std::endl; auto artifact_path = hnsw_build(dev_resources, hnsw_params, dataset_i8_view); - std::cout << "[stage 3] Deserialize layered HNSW index" << std::endl; - auto hnsw_index = hnsw_deserialize(dev_resources, hnsw_params, artifact_path, n_dim); + std::cout << "[stage 3] Materialize layered HNSW -> native hnswlib index" << std::endl; + hnsw_materialize( + dev_resources, hnsw_params, artifact_path, dataset_path, n_dim, native_index_path); + + std::cout << "[stage 4] Read materialized hnswlib index into memory" << std::endl; + auto hnsw_index = + hnsw_load_native(dev_resources, native_index_path, hnsw_params.metric, n_dim); - std::cout << "[stage 4] Search HNSW index" << std::endl; + std::cout << "[stage 5] Search HNSW index" << std::endl; hnsw_search(dev_resources, *hnsw_index, queries_i8_view); #else auto dataset_path = write_local_dataset(dataset_host_view, std::string{kBuildDir} + "/dataset.npy"); auto hnsw_params = make_hnsw_ace_params(kBuildDir, dataset_path); + const std::string native_index_path = std::string{kBuildDir} + "/hnsw_native.bin"; + std::cout << "[stage 2] Build layered HNSW index with ACE" << std::endl; auto artifact_path = hnsw_build(dev_resources, hnsw_params, dataset_host_view); - std::cout << "[stage 3] Deserialize layered HNSW index" << std::endl; - auto hnsw_index = hnsw_deserialize(dev_resources, hnsw_params, artifact_path, n_dim); + std::cout << "[stage 3] Materialize layered HNSW -> native hnswlib index" << std::endl; + hnsw_materialize( + dev_resources, hnsw_params, artifact_path, dataset_path, n_dim, native_index_path); + + std::cout << "[stage 4] Read materialized hnswlib index into memory" << std::endl; + auto hnsw_index = + hnsw_load_native(dev_resources, native_index_path, hnsw_params.metric, n_dim); - std::cout << "[stage 4] Search HNSW index" << std::endl; + std::cout << "[stage 5] Search HNSW index" << std::endl; hnsw_search(dev_resources, *hnsw_index, queries_host_view); #endif diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java index 3eef491b62..14d897521e 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java @@ -82,6 +82,44 @@ static HnswIndex build(CuVSResources resources, HnswIndexParams hnswParams, CuVS return CuVSProvider.provider().hnswIndexBuild(resources, hnswParams, dataset); } + /** + * Materializes a layered HNSW artifact into a standard hnswlib index file on + * disk. + * + * Materializes a {@code GPU_LAYERED_ON_DISK} artifact (graph topology only, + * stored in ACE order) plus a local dataset into a standard hnswlib index file, + * without ever holding the full materialized index in host memory. The + * resulting file is compatible with the original hnswlib library and can be read + * back with {@code hierarchy == CPU}. The element data type is read from the + * artifact header. + * + * @param resources The CuVS resources + * @param materializeParams Materialization parameters (dataset path, host-memory + * budget, threads) + * @param layeredArtifactPath Path to the layered HNSW artifact + * @param outputPath Path to the hnswlib index file to write + * @param dim The dimension of the vectors in the index + * @param metric The distance metric used to build the index + * @throws Throwable if an error occurs during materialization + */ + static void materializeToHnswlib( + CuVSResources resources, + HnswMaterializeParams materializeParams, + String layeredArtifactPath, + String outputPath, + int dim, + HnswIndexParams.CuvsDistanceType metric) + throws Throwable { + Objects.requireNonNull(resources); + Objects.requireNonNull(materializeParams); + Objects.requireNonNull(layeredArtifactPath); + Objects.requireNonNull(outputPath); + Objects.requireNonNull(metric); + CuVSProvider.provider() + .hnswMaterializeToHnswlib( + resources, materializeParams, layeredArtifactPath, outputPath, dim, metric); + } + /** * Builder helps configure and create an instance of {@link HnswIndex}. */ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java index 070cbedae1..6de60ed4f1 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java @@ -46,7 +46,16 @@ public enum CuvsHnswHierarchy { /** * Full hierarchy is built using the GPU */ - GPU(2); + GPU(2), + + /** + * GPU-built hierarchy stored as a layered on-disk topology artifact. + * + * The artifact stores graph topology only. When loading such an artifact, the + * {@code datasetPath} must point to the original-ID-ordered vectors used to + * reconstruct an in-memory HNSW index. + */ + GPU_LAYERED_ON_DISK(3); /** * The value for the enum choice. @@ -65,6 +74,8 @@ public static class HnswHierarchy { public static final CuvsHnswHierarchy NONE = CuvsHnswHierarchy.NONE; public static final CuvsHnswHierarchy CPU = CuvsHnswHierarchy.CPU; public static final CuvsHnswHierarchy GPU = CuvsHnswHierarchy.GPU; + public static final CuvsHnswHierarchy GPU_LAYERED_ON_DISK = + CuvsHnswHierarchy.GPU_LAYERED_ON_DISK; } private CuvsHnswHierarchy hierarchy = CuvsHnswHierarchy.GPU; @@ -74,6 +85,7 @@ public static class HnswHierarchy { private long m = 32; private CuvsDistanceType metric = CuvsDistanceType.L2Expanded; private HnswAceParams aceParams; + private String datasetPath; private HnswIndexParams( CuvsHnswHierarchy hierarchy, @@ -82,7 +94,8 @@ private HnswIndexParams( int vectorDimension, long m, CuvsDistanceType metric, - HnswAceParams aceParams) { + HnswAceParams aceParams, + String datasetPath) { this.hierarchy = hierarchy; this.efConstruction = efConstruction; this.numThreads = numThreads; @@ -90,6 +103,7 @@ private HnswIndexParams( this.m = m; this.metric = metric; this.aceParams = aceParams; + this.datasetPath = datasetPath; } /** @@ -153,6 +167,20 @@ public HnswAceParams getAceParams() { return aceParams; } + /** + * Gets the local dataset path used by layered HNSW deserialization. + * + * Required when {@code hierarchy == GPU_LAYERED_ON_DISK}: the artifact stores + * graph topology only, and loading reads the original-ID-ordered vectors from + * this path to reconstruct an in-memory HNSW index. Ignored for all other + * hierarchies. + * + * @return the dataset path, or null if not set + */ + public String getDatasetPath() { + return datasetPath; + } + @Override public String toString() { return "HnswIndexParams [hierarchy=" @@ -169,6 +197,8 @@ public String toString() { + metric + ", aceParams=" + aceParams + + ", datasetPath=" + + datasetPath + "]"; } @@ -184,6 +214,7 @@ public static class Builder { private long m = 32; private CuvsDistanceType metric = CuvsDistanceType.L2Expanded; private HnswAceParams aceParams; + private String datasetPath; /** * Constructs this Builder with an instance of Arena. @@ -276,6 +307,22 @@ public Builder withAceParams(HnswAceParams aceParams) { return this; } + /** + * Sets the local dataset path used by layered HNSW deserialization. + * + * Required when {@code hierarchy == GPU_LAYERED_ON_DISK}: the artifact stores + * graph topology only, and loading reads the original-ID-ordered vectors from + * this path to reconstruct an in-memory HNSW index. Ignored for all other + * hierarchies. + * + * @param datasetPath the local dataset path + * @return an instance of Builder + */ + public Builder withDatasetPath(String datasetPath) { + this.datasetPath = datasetPath; + return this; + } + /** * Builds an instance of {@link HnswIndexParams}. * @@ -289,7 +336,8 @@ public HnswIndexParams build() { vectorDimension, m, metric, - aceParams); + aceParams, + datasetPath); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswMaterializeParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswMaterializeParams.java new file mode 100644 index 0000000000..efbcc8acc1 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswMaterializeParams.java @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +/** + * Parameters for materializing a layered HNSW artifact into a standard hnswlib + * index file on disk. + * + * @since 26.08 + */ +public class HnswMaterializeParams { + + private final String datasetPath; + private final double maxHostMemoryGb; + private final int numThreads; + + private HnswMaterializeParams(String datasetPath, double maxHostMemoryGb, int numThreads) { + this.datasetPath = datasetPath; + this.maxHostMemoryGb = maxHostMemoryGb; + this.numThreads = numThreads; + } + + /** + * Gets the local dataset path holding the original-ID-ordered vectors used to + * build the artifact. + * + * @return the dataset path, or null if not set + */ + public String getDatasetPath() { + return datasetPath; + } + + /** + * Gets the upper bound on host memory (in GiB) used for the base-topology + * reorder buffer. When {@code <= 0}, the whole base topology is reordered in a + * single in-memory pass. + * + * @return the max host memory in GiB (0 means a single in-memory pass) + */ + public double getMaxHostMemoryGb() { + return maxHostMemoryGb; + } + + /** + * Gets the number of host threads to use. When 0, the maximum number of + * threads is used. + * + * @return the number of threads + */ + public int getNumThreads() { + return numThreads; + } + + @Override + public String toString() { + return "HnswMaterializeParams [datasetPath=" + + datasetPath + + ", maxHostMemoryGb=" + + maxHostMemoryGb + + ", numThreads=" + + numThreads + + "]"; + } + + /** + * Builder configures and creates an instance of {@link HnswMaterializeParams}. + */ + public static class Builder { + + private String datasetPath; + private double maxHostMemoryGb = 0; + private int numThreads = 0; + + /** + * Constructs this Builder. + */ + public Builder() {} + + /** + * Sets the local dataset path holding the original-ID-ordered vectors used to + * build the artifact. Supported formats match layered deserialization: + * {@code .npy} and ANN benchmark {@code *.bin} files with a + * {@code [uint32 rows, uint32 cols]} header ({@code .fbin}, {@code .f16bin}, + * {@code .u8bin}, {@code .i8bin}). + * + * @param datasetPath the local dataset path + * @return an instance of Builder + */ + public Builder withDatasetPath(String datasetPath) { + this.datasetPath = datasetPath; + return this; + } + + /** + * Sets the upper bound on host memory (in GiB) used for the base-topology + * reorder buffer. + * + * When {@code <= 0} (default), the whole base topology is reordered in a single + * in-memory pass (no temporary files). When set, the base topology is reordered + * through bucketed temporary files so that peak host memory stays close to this + * budget. + * + * @param maxHostMemoryGb the max host memory in GiB + * @return an instance of Builder + */ + public Builder withMaxHostMemoryGb(double maxHostMemoryGb) { + this.maxHostMemoryGb = maxHostMemoryGb; + return this; + } + + /** + * Sets the number of host threads to use. When 0 (default), the maximum number + * of threads is used. + * + * @param numThreads the number of threads + * @return an instance of Builder + */ + public Builder withNumThreads(int numThreads) { + this.numThreads = numThreads; + return this; + } + + /** + * Builds an instance of {@link HnswMaterializeParams}. + * + * @return an instance of {@link HnswMaterializeParams} + */ + public HnswMaterializeParams build() { + return new HnswMaterializeParams(datasetPath, maxHostMemoryGb, numThreads); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index c39578755c..9ba6a34475 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -135,6 +135,28 @@ HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) throws Throwable; + /** + * Materializes a layered HNSW artifact into a standard hnswlib index file on disk. + * + * @param resources The CuVS resources + * @param materializeParams Materialization parameters (dataset path, host-memory budget, threads) + * @param layeredArtifactPath Path to the layered HNSW artifact + * @param outputPath Path to the hnswlib index file to write + * @param dim The dimension of the vectors in the index + * @param metric The distance metric used to build the index + * @throws Throwable if an error occurs during materialization + */ + default void hnswMaterializeToHnswlib( + CuVSResources resources, + HnswMaterializeParams materializeParams, + String layeredArtifactPath, + String outputPath, + int dim, + HnswIndexParams.CuvsDistanceType metric) + throws Throwable { + throw new UnsupportedOperationException(); + } + /** Creates a new TieredIndex Builder. */ TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java index 950504bc5a..573699a39d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java @@ -160,6 +160,27 @@ public void close() { } } + static CloseableHandle createHnswMaterializeParamsNative() { + try (var localArena = Arena.ofConfined()) { + var paramsPtrPtr = localArena.allocate(cuvsHnswMaterializeParams_t); + checkCuVSError( + cuvsHnswMaterializeParamsCreate(paramsPtrPtr), "cuvsHnswMaterializeParamsCreate"); + var paramsPtr = paramsPtrPtr.get(cuvsHnswMaterializeParams_t, 0L); + return new CloseableHandle() { + @Override + public MemorySegment handle() { + return paramsPtr; + } + + @Override + public void close() { + checkCuVSError( + cuvsHnswMaterializeParamsDestroy(paramsPtr), "cuvsHnswMaterializeParamsDestroy"); + } + }; + } + } + static CloseableHandle createTieredIndexParams() { try (var localArena = Arena.ofConfined()) { var paramsPtrPtr = localArena.allocate(cuvsTieredIndexParams_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index ca528ac010..c6e241180b 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -6,6 +6,7 @@ import static com.nvidia.cuvs.internal.CuVSParamsHelper.createHnswAceParamsNative; import static com.nvidia.cuvs.internal.CuVSParamsHelper.createHnswIndexParams; +import static com.nvidia.cuvs.internal.CuVSParamsHelper.createHnswMaterializeParamsNative; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; @@ -19,6 +20,7 @@ import com.nvidia.cuvs.HnswAceParams; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.HnswIndexParams; +import com.nvidia.cuvs.HnswMaterializeParams; import com.nvidia.cuvs.HnswQuery; import com.nvidia.cuvs.HnswSearchParams; import com.nvidia.cuvs.SearchResults; @@ -27,6 +29,7 @@ import com.nvidia.cuvs.internal.panama.cuvsHnswAceParams; import com.nvidia.cuvs.internal.panama.cuvsHnswIndex; import com.nvidia.cuvs.internal.panama.cuvsHnswIndexParams; +import com.nvidia.cuvs.internal.panama.cuvsHnswMaterializeParams; import com.nvidia.cuvs.internal.panama.cuvsHnswSearchParams; import java.io.InputStream; import java.lang.foreign.Arena; @@ -197,7 +200,14 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { cuvsHnswIndex.dtype(indexReference.memorySegment, dtype); - try (var params = segmentFromIndexParams(hnswIndexParams); + // The metric only affects loading for the layered (GPU_LAYERED_ON_DISK) hierarchy, where the + // native side validates it against the artifact header. For the CPU/GPU/NONE hnswlib paths + // the metric is otherwise unused at load time. + int deserializeMetric = + hnswIndexParams.getHierarchy() == HnswIndexParams.CuvsHnswHierarchy.GPU_LAYERED_ON_DISK + ? hnswIndexParams.getMetric().value + : 0; + try (var params = segmentFromIndexParams(localArena, hnswIndexParams); var cuvsResourcesAccessor = resources.access()) { checkCuVSError( cuvsHnswDeserialize( @@ -205,7 +215,7 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { params.handle(), pathSeg, hnswIndexParams.getVectorDimension(), - 0, + deserializeMetric, indexReference.memorySegment), "cuvsHnswDeserialize"); } @@ -218,13 +228,21 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { } /** - * Allocates the configured search parameters in the MemorySegment. + * Allocates the configured index parameters in the MemorySegment. + * + * The {@code dataset_path} string (when set) is allocated from {@code arena}, so + * the caller must keep {@code arena} alive until the native call that consumes + * the params has returned. */ - private CloseableHandle segmentFromIndexParams(HnswIndexParams params) { + private CloseableHandle segmentFromIndexParams(Arena arena, HnswIndexParams params) { var hnswParams = createHnswIndexParams(); cuvsHnswIndexParams.hierarchy(hnswParams.handle(), params.getHierarchy().value); cuvsHnswIndexParams.ef_construction(hnswParams.handle(), params.getEfConstruction()); cuvsHnswIndexParams.num_threads(hnswParams.handle(), params.getNumThreads()); + String datasetPath = params.getDatasetPath(); + if (datasetPath != null) { + cuvsHnswIndexParams.dataset_path(hnswParams.handle(), arena.allocateFrom(datasetPath)); + } return hnswParams; } @@ -288,6 +306,56 @@ public static HnswIndex build(CuVSResources resources, HnswIndexParams hnswParam return new HnswIndexImpl(new IndexReference(hnswIndex), resources, hnswParams); } + /** + * Materializes a layered HNSW artifact into a standard hnswlib index file on disk. + * + * @param resources the CuVS resources + * @param materializeParams the materialization parameters + * @param layeredArtifactPath path to the layered HNSW artifact + * @param outputPath path to the hnswlib index file to write + * @param dim the dimension of the vectors in the index + * @param metric the distance metric used to build the index + * @throws Throwable if an error occurs during materialization + */ + public static void materializeToHnswlib( + CuVSResources resources, + HnswMaterializeParams materializeParams, + String layeredArtifactPath, + String outputPath, + int dim, + HnswIndexParams.CuvsDistanceType metric) + throws Throwable { + Objects.requireNonNull(resources); + Objects.requireNonNull(materializeParams); + Objects.requireNonNull(layeredArtifactPath); + Objects.requireNonNull(outputPath); + Objects.requireNonNull(metric); + + try (var localArena = Arena.ofConfined(); + var paramsHandle = createHnswMaterializeParamsNative()) { + MemorySegment paramsSeg = paramsHandle.handle(); + + String datasetPath = materializeParams.getDatasetPath(); + if (datasetPath != null) { + cuvsHnswMaterializeParams.dataset_path(paramsSeg, localArena.allocateFrom(datasetPath)); + } + cuvsHnswMaterializeParams.max_host_memory_gb( + paramsSeg, materializeParams.getMaxHostMemoryGb()); + cuvsHnswMaterializeParams.num_threads(paramsSeg, materializeParams.getNumThreads()); + + MemorySegment artifactSeg = buildMemorySegment(localArena, layeredArtifactPath); + MemorySegment outputSeg = buildMemorySegment(localArena, outputPath); + + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + int returnValue = + cuvsHnswMaterializeToHnswlib( + cuvsRes, paramsSeg, artifactSeg, outputSeg, dim, metric.value); + checkCuVSError(returnValue, "cuvsHnswMaterializeToHnswlib"); + } + } + } + private static CloseableHandle createHnswIndexParamsForBuild(Arena arena, HnswIndexParams params) { var hnswParams = createHnswIndexParams(); MemorySegment seg = hnswParams.handle(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 1d3199f26f..b7b03c8ce5 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -260,6 +260,19 @@ public HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswPar return HnswIndexImpl.build(resources, hnswParams, dataset); } + @Override + public void hnswMaterializeToHnswlib( + CuVSResources resources, + HnswMaterializeParams materializeParams, + String layeredArtifactPath, + String outputPath, + int dim, + HnswIndexParams.CuvsDistanceType metric) + throws Throwable { + HnswIndexImpl.materializeToHnswlib( + resources, materializeParams, layeredArtifactPath, outputPath, dim, metric); + } + @Override public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { return TieredIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); diff --git a/python/cuvs/cuvs/neighbors/hnsw/__init__.py b/python/cuvs/cuvs/neighbors/hnsw/__init__.py index f91835b7c5..3e8541b603 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/__init__.py +++ b/python/cuvs/cuvs/neighbors/hnsw/__init__.py @@ -7,11 +7,13 @@ ExtendParams, Index, IndexParams, + MaterializeParams, SearchParams, build, extend, from_cagra, load, + materialize_to_hnswlib, save, search, ) @@ -21,10 +23,12 @@ "IndexParams", "Index", "ExtendParams", + "MaterializeParams", "build", "extend", "SearchParams", "load", + "materialize_to_hnswlib", "save", "search", "from_cagra", diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd index 9ffb295ad3..1dff7a0809 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd @@ -20,6 +20,7 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil: NONE CPU GPU + GPU_LAYERED_ON_DISK ctypedef struct cuvsHnswAceParams: size_t npartitions @@ -41,6 +42,7 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil: size_t M cuvsDistanceType metric cuvsHnswAceParams_t ace_params + const char* dataset_path ctypedef cuvsHnswIndexParams* cuvsHnswIndexParams_t @@ -105,3 +107,24 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil: int32_t dim, cuvsDistanceType metric, cuvsHnswIndex_t index) except + + + ctypedef struct cuvsHnswMaterializeParams: + const char* dataset_path + double max_host_memory_gb + int num_threads + + ctypedef cuvsHnswMaterializeParams* cuvsHnswMaterializeParams_t + + cuvsError_t cuvsHnswMaterializeParamsCreate( + cuvsHnswMaterializeParams_t* params) + + cuvsError_t cuvsHnswMaterializeParamsDestroy( + cuvsHnswMaterializeParams_t params) + + cuvsError_t cuvsHnswMaterializeToHnswlib( + cuvsResources_t res, + cuvsHnswMaterializeParams_t params, + const char* layered_artifact_path, + const char* output_path, + int32_t dim, + cuvsDistanceType metric) except + diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index 6757695de3..bff11001d0 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -124,10 +124,15 @@ cdef class IndexParams: Parameters ---------- hierarchy : string, default = "gpu" (optional) - The hierarchy of the HNSW index. Valid values are ["none", "cpu", "gpu"]. + The hierarchy of the HNSW index. Valid values are + ["none", "cpu", "gpu", "gpu_layered_on_disk"]. - "none": No hierarchy is built. - "cpu": Hierarchy is built using CPU. - "gpu": Hierarchy is built using GPU. + - "gpu_layered_on_disk": The index artifact stores graph topology only + (built on the GPU and stored as a layered on-disk artifact). When + loading such an artifact with `load()`, `dataset_path` must point to + the original-ID-ordered vectors used to reconstruct the index. ef_construction : int, default = 200 (optional) Maximum number of candidate list size used during construction when hierarchy is `cpu`. @@ -148,14 +153,22 @@ cdef class IndexParams: ace_params : AceParams, default = None (optional) ACE parameters for building HNSW index using ACE algorithm. If set, enables the build() function to use ACE for index construction. + dataset_path : string, default = None (optional) + Local dataset path used by layered HNSW deserialization. Required when + `hierarchy == "gpu_layered_on_disk"`: the artifact stores graph + topology only, and `load()` reads the original-ID-ordered vectors from + this path to reconstruct an in-memory HNSW index. Ignored for all other + hierarchies. """ cdef cuvsHnswIndexParams* params cdef AceParams _ace_params + cdef object _dataset_path_bytes def __cinit__(self): check_cuvs(cuvsHnswIndexParamsCreate(&self.params)) self._ace_params = None + self._dataset_path_bytes = None def __dealloc__(self): check_cuvs(cuvsHnswIndexParamsDestroy(self.params)) @@ -166,16 +179,20 @@ cdef class IndexParams: num_threads=0, M=32, metric="sqeuclidean", - ace_params=None): + ace_params=None, + dataset_path=None): if hierarchy == "none": self.params.hierarchy = cuvsHnswHierarchy.NONE elif hierarchy == "cpu": self.params.hierarchy = cuvsHnswHierarchy.CPU elif hierarchy == "gpu": self.params.hierarchy = cuvsHnswHierarchy.GPU + elif hierarchy == "gpu_layered_on_disk": + self.params.hierarchy = cuvsHnswHierarchy.GPU_LAYERED_ON_DISK else: raise ValueError("Invalid hierarchy type." - " Valid values are 'none', 'cpu', and 'gpu'.") + " Valid values are 'none', 'cpu', 'gpu', and" + " 'gpu_layered_on_disk'.") self.params.ef_construction = ef_construction self.params.num_threads = num_threads self.params.M = M @@ -189,6 +206,12 @@ cdef class IndexParams: else: self.params.ace_params = NULL + if dataset_path is not None: + self._dataset_path_bytes = dataset_path.encode('utf-8') + self.params.dataset_path = self._dataset_path_bytes + else: + self.params.dataset_path = NULL + @property def hierarchy(self): if self.params.hierarchy == cuvsHnswHierarchy.NONE: @@ -197,6 +220,8 @@ cdef class IndexParams: return "cpu" elif self.params.hierarchy == cuvsHnswHierarchy.GPU: return "gpu" + elif self.params.hierarchy == cuvsHnswHierarchy.GPU_LAYERED_ON_DISK: + return "gpu_layered_on_disk" @property def ef_construction(self): @@ -214,6 +239,12 @@ cdef class IndexParams: def ace_params(self): return self._ace_params + @property + def dataset_path(self): + if self.params.dataset_path is not NULL: + return self.params.dataset_path.decode('utf-8') + return None + cdef class Index: """ @@ -270,6 +301,68 @@ cdef class ExtendParams: return self.params.num_threads +cdef class MaterializeParams: + """ + Parameters for materializing a layered HNSW artifact into an hnswlib + index on disk. + + Parameters + ---------- + dataset_path : string, default = None (optional) + Local dataset path holding the original-ID-ordered vectors used to + build the artifact. Supported formats match layered deserialization: + `.npy` and ANN benchmark `*.bin` files with a + `[uint32 rows, uint32 cols]` header (`.fbin`, `.f16bin`, `.u8bin`, + `.i8bin`). + max_host_memory_gb : float, default = 0 (optional) + Upper bound on host memory (in GiB) used for the base-topology reorder + buffer. When <= 0, the whole base topology is reordered in a single + in-memory pass (no temporary files). When set, the base topology is + reordered through bucketed temporary files so that peak host memory + stays close to this budget. + num_threads : int, default = 0 (optional) + Number of host threads to use. When 0, the maximum number of threads + is used. + """ + + cdef cuvsHnswMaterializeParams* params + cdef object _dataset_path_bytes + + def __cinit__(self): + check_cuvs(cuvsHnswMaterializeParamsCreate(&self.params)) + self._dataset_path_bytes = None + + def __dealloc__(self): + if self.params is not NULL: + check_cuvs(cuvsHnswMaterializeParamsDestroy(self.params)) + + def __init__(self, *, + dataset_path=None, + max_host_memory_gb=0, + num_threads=0): + if dataset_path is not None: + self._dataset_path_bytes = dataset_path.encode('utf-8') + self.params.dataset_path = self._dataset_path_bytes + else: + self.params.dataset_path = NULL + self.params.max_host_memory_gb = max_host_memory_gb + self.params.num_threads = num_threads + + @property + def dataset_path(self): + if self.params.dataset_path is not NULL: + return self.params.dataset_path.decode('utf-8') + return None + + @property + def max_host_memory_gb(self): + return self.params.max_host_memory_gb + + @property + def num_threads(self): + return self.params.num_threads + + @auto_sync_resources def save(filename, Index index, resources=None): """ @@ -406,6 +499,84 @@ def load(IndexParams index_params, filename, dim, dtype, metric="sqeuclidean", return idx +@auto_sync_resources +def materialize_to_hnswlib(MaterializeParams materialize_params, + layered_artifact_path, + output_path, + dim, + metric="sqeuclidean", + resources=None): + """ + Materialize a layered HNSW artifact into a standard hnswlib index file + on disk. + + Materializes a `gpu_layered_on_disk` artifact (graph topology only, stored + in ACE order) plus a local dataset into a standard hnswlib index file, + without ever holding the full materialized index in host memory. The + resulting file is compatible with the original hnswlib library and can be + read back through `load()` with `hierarchy="cpu"`. The element data type + (float32, float16, uint8, int8) is read from the artifact header. + + Parameters + ---------- + materialize_params : MaterializeParams + Materialization parameters. `dataset_path` must point to the + original-ID-ordered vectors used to build the artifact. + layered_artifact_path : string + Path to the layered HNSW artifact. + output_path : string + Path to the hnswlib index file to write. + dim : int + Dimensions of the training dataset. + metric : string denoting the metric type, default="sqeuclidean" + Valid values for metric: ["sqeuclidean", "inner_product"], where + - sqeuclidean is the euclidean distance without the square root + operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, + - inner_product distance is defined as + distance(a, b) = \\sum_i a_i * b_i. + {resources_docstring} + + Examples + -------- + >>> import numpy as np + >>> from cuvs.neighbors import hnsw + >>> n_features = 50 + >>> # Assume a layered artifact was produced by an ACE GPU build and the + >>> # original-ID-ordered vectors are stored in "dataset.fbin". + >>> materialize_params = hnsw.MaterializeParams( + ... dataset_path="dataset.fbin" + ... ) + >>> hnsw.materialize_to_hnswlib( + ... materialize_params, + ... "layered_artifact.cuvs", + ... "index.bin", + ... n_features, + ... metric="sqeuclidean", + ... ) + >>> # The materialized index can be loaded as a standard hnswlib index. + >>> index = hnsw.load( + ... hnsw.IndexParams(hierarchy="cpu"), + ... "index.bin", + ... n_features, + ... np.float32, + ... "sqeuclidean", + ... ) + """ + cdef string c_artifact = layered_artifact_path.encode('utf-8') + cdef string c_output = output_path.encode('utf-8') + cdef cuvsDistanceType distance_type = DISTANCE_TYPES[metric] + cdef cuvsResources_t res = resources.get_c_obj() + + check_cuvs(cuvsHnswMaterializeToHnswlib( + res, + materialize_params.params, + c_artifact.c_str(), + c_output.c_str(), + dim, + distance_type + )) + + @auto_sync_resources def from_cagra(IndexParams index_params, cagra.Index cagra_index, temporary_index_path=None, resources=None): From 2d86facb5b4033da6a7c93f31c1d1f9449af81fd Mon Sep 17 00:00:00 2001 From: Micka Date: Tue, 22 Sep 2026 09:26:50 +0200 Subject: [PATCH 10/12] Add support for BBQ compressed datasets in NN-Descent and CAGRA (#2654) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BBQ is a data compression technique that is a variant of RaBitQ. It is used in Lucene and Elasticsearch to compresses vectors while keeping search fast and accurate. Supporting it in cuVS enables indexing on data compressed up to ~32× With this PR cuVS users will be able to build a CAGRA index using the NN-Descent algorithm with their BBQ-Quantized Data. Closes #2326 Authors: - Micka (https://github.com/lowener) - Artem M. Chirkin (https://github.com/achirkin) - Yuxi Sun (https://github.com/sherylll) - Anupam (https://github.com/aamijar) - Igor Motov (https://github.com/imotov) Approvers: - Tamas Bela Feher (https://github.com/tfeher) - Kyle Edwards (https://github.com/KyleFromNVIDIA) - Igor Motov (https://github.com/imotov) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/NVIDIA/cuvs/pull/2654 --- c/CMakeLists.txt | 1 + c/include/cuvs/core/all.h | 1 + c/include/cuvs/core/dataset.h | 3 +- c/include/cuvs/neighbors/cagra.h | 17 +- c/include/cuvs/preprocessing/quantize/bbq.h | 142 ++ c/src/neighbors/cagra.cpp | 180 ++- c/src/preprocessing/quantize/bbq.cpp | 243 +++ c/tests/core/dataset_c.cu | 131 ++ c/tests/neighbors/ann_cagra_c.cu | 175 ++ ci/validate_wheel.sh | 4 +- cpp/CMakeLists.txt | 4 +- cpp/include/cuvs/neighbors/cagra.hpp | 159 ++ cpp/include/cuvs/neighbors/common.hpp | 1 + cpp/include/cuvs/neighbors/nn_descent.hpp | 55 +- .../cuvs/preprocessing/quantize/bbq.hpp | 388 +++++ cpp/internal/CMakeLists.txt | 9 +- .../preprocessing/bbq_cpu_quantize.hpp | 629 ++++++++ cpp/src/neighbors/cagra.cuh | 8 +- cpp/src/neighbors/cagra_build_inst.cu.in | 8 + cpp/src/neighbors/cagra_search_inst.cu.in | 3 + cpp/src/neighbors/cagra_serialize.cuh | 29 + .../neighbors/detail/cagra/cagra_build.cuh | 120 ++ .../detail/cagra/cagra_serialize.cuh | 31 +- cpp/src/neighbors/detail/cagra/factory.cuh | 22 +- .../neighbors/detail/cagra/graph_shared.cu | 205 +++ .../neighbors/detail/cagra/graph_shared.cuh | 16 + cpp/src/neighbors/detail/nn_descent.cuh | 1411 ++++++++++++++++- cpp/src/neighbors/detail/nn_descent_gnnd.hpp | 15 +- cpp/src/neighbors/nn_descent.cuh | 19 +- cpp/src/neighbors/nn_descent_gnnd_inst.cu | 8 +- cpp/src/neighbors/nn_descent_inst.cu.in | 19 +- cpp/src/preprocessing/quantize/bbq.cu | 48 + .../quantize/detail/bbq_distance.cuh | 515 ++++++ cpp/tests/CMakeLists.txt | 9 + .../neighbors/ann_cagra/test_bbq_uint32_t.cu | 20 + cpp/tests/neighbors/ann_cagra_bbq.cuh | 379 +++++ .../ann_nn_descent/test_bbq_uint32_t.cu | 14 + cpp/tests/neighbors/ann_nn_descent_bbq.cuh | 208 +++ cpp/tests/neighbors/ann_utils.cuh | 7 + fern/docs.yml | 6 + fern/pages/c_api/c-api-core-dataset.md | 4 +- fern/pages/c_api/c-api-neighbors-cagra.md | 8 +- .../c_api/c-api-preprocessing-quantize-bbq.md | 153 ++ fern/pages/c_api/index.md | 1 + fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 98 ++ .../cpp_api/cpp-api-neighbors-nn-descent.md | 108 ++ .../cpp-api-preprocessing-quantize-bbq.md | 153 ++ fern/pages/cpp_api/index.md | 1 + fern/pages/java_api/index.md | 1 + .../java-api-com-nvidia-cuvs-bbqquantizer.md | 31 + .../java-api-com-nvidia-cuvs-cagraindex.md | 22 +- .../java/com/nvidia/cuvs/BbqQuantizer.java | 195 +++ .../main/java/com/nvidia/cuvs/CagraIndex.java | 13 + .../nvidia/cuvs/internal/CagraIndexImpl.java | 405 ++++- .../test/java/com/nvidia/cuvs/CagraBbqIT.java | 673 ++++++++ 55 files changed, 6939 insertions(+), 189 deletions(-) create mode 100644 c/include/cuvs/preprocessing/quantize/bbq.h create mode 100644 c/src/preprocessing/quantize/bbq.cpp create mode 100644 cpp/include/cuvs/preprocessing/quantize/bbq.hpp create mode 100644 cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp create mode 100644 cpp/src/preprocessing/quantize/bbq.cu create mode 100644 cpp/src/preprocessing/quantize/detail/bbq_distance.cuh create mode 100644 cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu create mode 100644 cpp/tests/neighbors/ann_cagra_bbq.cuh create mode 100644 cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu create mode 100644 cpp/tests/neighbors/ann_nn_descent_bbq.cuh create mode 100644 fern/pages/c_api/c-api-preprocessing-quantize-bbq.md create mode 100644 fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md create mode 100644 fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index eab7560dfc..59932c32ee 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -103,6 +103,7 @@ add_library( src/neighbors/all_neighbors.cpp src/preprocessing/pca.cpp src/preprocessing/quantize/binary.cpp + src/preprocessing/quantize/bbq.cpp src/preprocessing/quantize/pq.cpp src/preprocessing/quantize/scalar.cpp src/distance/pairwise_distance.cpp diff --git a/c/include/cuvs/core/all.h b/c/include/cuvs/core/all.h index 1b693a355e..80b95992eb 100644 --- a/c/include/cuvs/core/all.h +++ b/c/include/cuvs/core/all.h @@ -44,6 +44,7 @@ #include #include +#include #include #include diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h index 0fbc168759..1a1ab2565b 100644 --- a/c/include/cuvs/core/dataset.h +++ b/c/include/cuvs/core/dataset.h @@ -21,7 +21,8 @@ extern "C" { typedef enum { CUVS_DATASET_LAYOUT_STANDARD = 0, CUVS_DATASET_LAYOUT_PADDED = 1, - CUVS_DATASET_LAYOUT_PQ = 2 + CUVS_DATASET_LAYOUT_PQ = 2, + CUVS_DATASET_LAYOUT_BBQ = 3 } cuvsDatasetLayout_t; /** diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 6a8b854fc2..602f52dbfc 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -654,7 +654,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * * The memory space and layout \p dataset was constructed with select the C++ build overload. * Build the handle with an owning factory or the matching dataset view factory - * (`cuvsDatasetMakePaddedView` / `cuvsDatasetMakeStandardView`). + * (`cuvsDatasetMakePaddedView`, `cuvsDatasetMakeStandardView`, or + * `cuvsDatasetMakeBbqView`). * * Note that a dataset residing in host memory produces a host-backed index, which * must be made search-ready with `cuvsCagraUpdateDataset` (using a device-padded @@ -696,6 +697,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q * index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the * VPQ dataset must outlive the index because the index stores a non-owning view. + * A `CUVS_DATASET_LAYOUT_BBQ` dataset builds a graph-only index; attach a searchable dataset with + * `cuvsCagraUpdateDataset` before search. * * @param[in] res cuvsResources_t opaque C handle * @param[in] params cuvsCagraIndexParams_t used to build CAGRA index @@ -860,8 +863,9 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res, /** * Save the CAGRA graph to file without its dataset. * - * This supports dense and PQ-backed indexes. The dataset must be attached separately after loading - * the graph. + * This supports dense, PQ-backed, and BBQ-built indexes. The serialized file does not contain + * vector data. After deserialization the index cannot be searched until a compatible dataset is + * attached with `cuvsCagraUpdateDataset`. * * Experimental, both the API and the serialization format are subject to change. * @@ -876,9 +880,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraph(cuvsResources_t res, /** * Save the CAGRA graph and its attached dataset to file. * - * The index stores a non-owning dataset view. The caller must keep the dataset backing that view - * alive while this function runs. Returns CUVS_ERROR without modifying the destination file if - * the index has no attached dataset. PQ datasets are not serialized by this function. + * The index stores a non-owning dataset view. The caller must keep the memory of the dataset + * backing that view alive while this function runs. Returns CUVS_ERROR without modifying the + * destination file if the index has no attached dataset. PQ and BBQ datasets are not serialized + * by this function. * * Experimental, both the API and the serialization format are subject to change. * diff --git a/c/include/cuvs/preprocessing/quantize/bbq.h b/c/include/cuvs/preprocessing/quantize/bbq.h new file mode 100644 index 0000000000..b08a989a18 --- /dev/null +++ b/c/include/cuvs/preprocessing/quantize/bbq.h @@ -0,0 +1,142 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup preprocessing_c_bbq C API for Better Binary Quantization datasets + * @{ + */ + +/** + * Storage layout of BBQ/OSQ quantized component codes in each dataset row. + * CUVS_BBQ_CODE_LAYOUT_PACKED_1B: Each dimension is quantized to a single bit and packed into bytes. Reflects + * Lucene's OptimizedScalarQuantizer.packAsBinary. + * CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B: Each dimension is quantized to 2 bits, stored as 2 bitplanes. + * Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only + * (paired with a transposed_4b or packed_1b operand); + * CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B: Each dimension is quantized to 4 bits, optimized for bitwise operations. + * Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of + * every dimension is in the first set dimensions bits, or (dimensions/8) + * bytes. The second, third, and fourth bits are in the second, third, and + * fourth set of dimensions bits, respectively. Format used for queries. + * CUVS_BBQ_CODE_LAYOUT_PACKED_4B: Each dimension is quantized to 4 bits, two values are packed into each output + * byte. + * CUVS_BBQ_CODE_LAYOUT_PACKED_7B: Each dimension is quantized to 7 bits and treated as a signed value. + * CUVS_BBQ_CODE_LAYOUT_PACKED_8B: Each dimension is quantized to 8 bits and treated as an unsigned value. + */ +typedef enum { + CUVS_BBQ_CODE_LAYOUT_PACKED_1B = 0, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_7B, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B +} cuvsBbqCodeLayout_t; + + +/** + * @brief Better Binary Quantization + * ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) + * is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas + * introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et + * al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective + * factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements + * this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch + * exposes it as BBQ. + * + * BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared + * centroid, no random rotation, and OSQ codes. + * + * RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain + * additional per-vector information so search is better than naïve sign-bit comparison. They differ + * in transformation and scale representation. RaBitQ commonly separates residual magnitude from + * direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector + * scalar intervals to interpret the compressed residual codes. + */ +typedef struct cuvsBbqQuantizer { + uintptr_t addr; + void (*destroy_addr)(void*); + DLDataType dtype; + bool is_owning; +} cuvsBbqQuantizer; +typedef cuvsBbqQuantizer* cuvsBbqQuantizer_t; + +/** + * @brief Create a BBQ quantizer view from caller-owned device tensors. + * + * Tensors are not copied and must remain valid while a derived dataset is in use. + * + * @param[in] codes uint8 device matrix containing encoded rows + * @param[in] lower_intervals float32 device vector with one lower interval per row + * @param[in] upper_intervals float32 device vector with one upper interval per row + * @param[in] additional_corrections float32 device vector with one correction per row + * @param[in] quantized_component_sums int32 device vector with one component sum per row + * @param[in] centroid device vector containing the dataset centroid + * @param[in] dequant_delta float32 device vector with one dequantization delta per row + * @param[in] dequant_sum_delta float32 device vector with one delta-times-sum value per row + * @param[in] row_norm float32 device vector with one original-space squared norm per row + * @param[in] layout encoded code layout + * @param[in] metric distance metric associated with the encoded dataset + * @param[in] centroid_norm_sq squared norm of the centroid + * @param[out] quantizer newly allocated non-owning quantizer handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsBbqQuantizerCreateView( + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq, + cuvsBbqQuantizer_t* quantizer); + +/** + * @brief Destroy a BBQ quantizer without destroying its caller-owned tensors. + * + * @param[in] quantizer quantizer handle to destroy + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer); + +/** + * @brief Create a non-owning device BBQ dataset view. + * + * Accepts one symmetric quantizer or two compatible asymmetric quantizers. + * + * @param[in] res cuVS resources + * @param[in] quantizers array containing one or two BBQ quantizer handles + * @param[in] num_quantizers number of elements in `quantizers` + * @param[out] dataset newly allocated non-owning BBQ dataset handle + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t res, + cuvsBbqQuantizer_t* quantizers, + size_t num_quantizers, + cuvsDataset_t* dataset); + +/** @} */ + +#ifdef __cplusplus +} +#endif diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 57e8168c1e..97cc9d8bac 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -24,12 +24,14 @@ #include "../core/exceptions.hpp" #include "../core/interop.hpp" #include +#include #include #include #include #include #include #include +#include #include #include "../core/exceptions.hpp" #include "../core/interop.hpp" @@ -57,6 +59,7 @@ struct sg_cagra_c_api_index_box { device_padded, device_standard, device_vpq, + device_bbq, host_padded, host_standard } layout; @@ -72,6 +75,8 @@ constexpr auto sg_cagra_index_layout_from_view() return sg_cagra_c_api_index_box::dataset_layout::device_padded; } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_vpq; + } else if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_bbq; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::host_standard; } else { @@ -79,7 +84,11 @@ constexpr auto sg_cagra_index_layout_from_view() } } -template +template static void with_index_by_layout(sg_cagra_c_api_index_box* box, const char* null_handle_err, const char* host_not_allowed_err, @@ -116,6 +125,16 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, fn(*idx); break; } + case sg_cagra_c_api_index_box::dataset_layout::device_bbq: { + if constexpr (AllowBbq) { + using index_t = cuvs::neighbors::cagra::device_bbq_index; + auto* idx = reinterpret_cast(box->index_ptr); + fn(*idx); + } else { + RAFT_FAIL("This operation does not support a BBQ-backed index"); + } + break; + } case sg_cagra_c_api_index_box::dataset_layout::host_standard: { if constexpr (AllowHost) { auto* idx = @@ -390,6 +409,92 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) } } +template +void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view const& dataset) +{ + using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + if (dataset.quantizers.size() == 1) { + auto layout = dataset.quantizers.front().layout; + RAFT_EXPECTS(layout == layout_t::packed_1b || layout == layout_t::transposed_2b || + layout == layout_t::packed_4b || layout == layout_t::packed_7b || + layout == layout_t::packed_8b, + "cuvsDatasetMakeBbqView: transposed_4b requires an asymmetric layout pair"); + return; + } + + const bool has_1b = dataset.has_layout(layout_t::packed_1b); + const bool has_2bt = dataset.has_layout(layout_t::transposed_2b); + const bool has_4b = dataset.has_layout(layout_t::packed_4b); + const bool has_4bt = dataset.has_layout(layout_t::transposed_4b); + RAFT_EXPECTS((has_1b && (has_4b || has_2bt || has_4bt)) || (has_2bt && has_4bt), + "cuvsDatasetMakeBbqView: unsupported asymmetric BBQ layout pair"); +} + +template +auto get_cpp_bbq_quantizer_view(cuvsBbqQuantizer_t quantizer) + -> cuvs::preprocessing::quantize::bbq::quantizer_view +{ + using owner_t = cuvs::preprocessing::quantize::bbq::quantizer; + using view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + RAFT_EXPECTS(quantizer != nullptr && quantizer->addr != 0, + "cuvsDatasetMakeBbqView: null quantizer storage"); + if (quantizer->is_owning) { + return reinterpret_cast(quantizer->addr)->view(); + } + return *reinterpret_cast(quantizer->addr); +} + +template +auto make_bbq_dataset_view(cuvsBbqQuantizer_t* quantizers, std::size_t num_quantizers) + -> std::unique_ptr> +{ + using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + auto dataset = std::make_unique(); + int64_t expected_rows{-1}; + uint32_t expected_dim{}; + cuvs::distance::DistanceType expected_metric{}; + + for (std::size_t i = 0; i < num_quantizers; ++i) { + RAFT_EXPECTS(quantizers[i] != nullptr, + "cuvsDatasetMakeBbqView: quantizer %zu is null", + i); + auto quantizer = get_cpp_bbq_quantizer_view(quantizers[i]); + if (i == 0) { + expected_rows = quantizer.n_rows(); + expected_dim = quantizer.dim(); + expected_metric = quantizer.metric; + } else { + RAFT_EXPECTS(quantizer.n_rows() == expected_rows && quantizer.dim() == expected_dim, + "cuvsDatasetMakeBbqView: all quantizers must have the same shape"); + RAFT_EXPECTS(quantizer.metric == expected_metric, + "cuvsDatasetMakeBbqView: all quantizers must use the same metric"); + } + dataset->add_quantizer(quantizer); + } + + return dataset; +} + +template +void make_and_bind_bbq_dataset(cuvsBbqQuantizer_t* quantizers, + std::size_t num_quantizers, + DLDataType dtype, + cuvsDataset_t* output) +{ + using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + auto view = make_bbq_dataset_view(quantizers, num_quantizers); + validate_bbq_layouts(*view); + + auto handle = std::make_unique(); + handle->addr = reinterpret_cast(view.release()); + handle->destroy_addr = &destroy_typed_addr; + handle->dtype = dtype; + handle->mem_type = CUVS_DATASET_MEM_TYPE_DEVICE; + handle->layout = CUVS_DATASET_LAYOUT_BBQ; + handle->is_owning = false; + *output = handle.release(); +} + using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset; using device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; @@ -845,7 +950,7 @@ void _search(cuvsResources_t res, { auto res_ptr = reinterpret_cast(res); auto* box = reinterpret_cast(index.addr); - with_index_by_layout( + with_index_by_layout( box, "cuvsCagraSearch: null index handle", "cuvsCagraSearch: host index must be converted to device first via " @@ -1014,7 +1119,11 @@ void _serialize(cuvsResources_t res, const char *filename, with_index_by_layout(box, null_handle_err, "", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + if constexpr (cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_EXPECTS(!include_dataset, + "cuvsCagraSerializeGraphAndDataset is not supported for BBQ indices"); + cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx); + } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { RAFT_EXPECTS( !include_dataset, "cuvsCagraSerializeGraphAndDataset is not supported for PQ indices; serialize the PQ " @@ -1186,8 +1295,9 @@ void _serialize_to_hnswlib(cuvsResources_t res, const char *filename, "cuvsCagraSerializeToHnswlib: host indices are allowed", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for VPQ indices"); + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for quantized dataset layouts"); } else { cuvs::neighbors::cagra::serialize_to_hnswlib( *res_ptr, std::string(filename), idx); @@ -1268,8 +1378,10 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) "cuvsCagraIndexGetDataset: host indices are allowed", [&](auto& idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cuvsCagraIndexGetDataset does not expose VPQ datasets as dense DLPack tensors"); + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v) { + RAFT_FAIL( + "cuvsCagraIndexGetDataset does not expose quantized datasets as dense DLPack tensors"); } else { cuvs::core::to_dlpack(idx.dataset().view(), dataset); } @@ -1654,6 +1766,46 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t, + cuvsBbqQuantizer_t* quantizers, + size_t num_quantizers, + cuvsDataset_t* dataset) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(dataset != nullptr, "cuvsDatasetMakeBbqView: null output dataset"); + *dataset = nullptr; + RAFT_EXPECTS(quantizers != nullptr, "cuvsDatasetMakeBbqView: null quantizer array"); + RAFT_EXPECTS(num_quantizers == 1 || num_quantizers == 2, + "cuvsDatasetMakeBbqView: expected one or two quantizers"); + RAFT_EXPECTS(quantizers[0] != nullptr && quantizers[0]->addr != 0, + "cuvsDatasetMakeBbqView: first quantizer has null storage"); + + auto dtype = quantizers[0]->dtype; + for (size_t i = 1; i < num_quantizers; ++i) { + RAFT_EXPECTS(quantizers[i] != nullptr && quantizers[i]->addr != 0, + "cuvsDatasetMakeBbqView: quantizer %zu has null storage", + i); + RAFT_EXPECTS(quantizers[i]->dtype.code == dtype.code && + quantizers[i]->dtype.bits == dtype.bits && + quantizers[i]->dtype.lanes == dtype.lanes, + "cuvsDatasetMakeBbqView: all quantizers must use the same dtype"); + } + if (dtype.code == kDLFloat && dtype.bits == 32) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + make_and_bind_bbq_dataset(quantizers, num_quantizers, dtype, dataset); + } else { + RAFT_FAIL("cuvsDatasetMakeBbqView: unsupported centroid dtype: code=%d, bits=%d", + dtype.code, + dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsDatasetMakePQ(cuvsResources_t res, cuvsPqParams_t params, cuvsDataset_t dataset, @@ -1766,6 +1918,17 @@ static void build_dispatch_on_mem_type_and_layout(raft::resources* res_ptr, cuvsDataset_t dataset, cuvsCagraIndex_t index) { + if (dataset->layout == CUVS_DATASET_LAYOUT_BBQ) { + RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE, + "cuvsCagraBuild: BBQ dataset must be device-resident"); + using owner_t = cuvs::neighbors::device_bbq_dataset; + using view_t = cuvs::neighbors::device_bbq_dataset_view; + with_dataset_view(dataset, [&](auto const& view) { + build_index_from_dataset_view(res_ptr, params, view, index); + }); + return; + } + auto const is_padded = dataset->layout == CUVS_DATASET_LAYOUT_PADDED; if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) { @@ -1925,6 +2088,9 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, auto index = *index_c_ptr; auto* box = reinterpret_cast(index.addr); RAFT_EXPECTS(box != nullptr, "cuvsCagraSearch: null index handle"); + RAFT_EXPECTS(box->layout != sg_cagra_c_api_index_box::dataset_layout::device_bbq, + "cuvsCagraSearch: attach a device padded or PQ dataset to the BBQ-built index " + "with cuvsCagraUpdateDataset before searching"); RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded || box->layout == sg_cagra_c_api_index_box::dataset_layout::device_vpq, "cuvsCagraSearch: index must be device-padded or VPQ. For standard indices, call " diff --git a/c/src/preprocessing/quantize/bbq.cpp b/c/src/preprocessing/quantize/bbq.cpp new file mode 100644 index 0000000000..a2415c661d --- /dev/null +++ b/c/src/preprocessing/quantize/bbq.cpp @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include "../../core/exceptions.hpp" +#include "../../core/interop.hpp" + +#include +#include + +#include +#include + +namespace { + +template +void destroy_typed_addr(void* ptr) +{ + delete reinterpret_cast(ptr); +} + +auto to_cpp_layout(cuvsBbqCodeLayout_t layout) + -> cuvs::preprocessing::quantize::bbq::bbq_code_layout +{ + using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + switch (layout) { + case CUVS_BBQ_CODE_LAYOUT_PACKED_1B: return layout_t::packed_1b; + case CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B: return layout_t::transposed_2b; + case CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B: return layout_t::transposed_4b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_4B: return layout_t::packed_4b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_7B: return layout_t::packed_7b; + case CUVS_BBQ_CODE_LAYOUT_PACKED_8B: return layout_t::packed_8b; + } + RAFT_FAIL("cuvsBbqQuantizerCreateView: invalid BBQ code layout"); +} + +template +auto make_cpp_quantizer_view( + DLManagedTensor* codes_tensor, + DLManagedTensor* lower_intervals_tensor, + DLManagedTensor* upper_intervals_tensor, + DLManagedTensor* additional_corrections_tensor, + DLManagedTensor* quantized_component_sums_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* dequant_delta_tensor, + DLManagedTensor* dequant_sum_delta_tensor, + DLManagedTensor* row_norm_tensor, + cuvsBbqCodeLayout_t c_layout, + cuvsDistanceType metric, + float centroid_norm_sq) + -> cuvs::preprocessing::quantize::bbq::quantizer_view +{ + using quantizer_view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + using codes_view_t = raft::device_matrix_view; + using float_view_t = raft::device_vector_view; + using int_view_t = raft::device_vector_view; + using centroid_view_t = raft::device_vector_view; + + auto codes = cuvs::core::from_dlpack(codes_tensor); + auto lower = cuvs::core::from_dlpack(lower_intervals_tensor); + auto upper = cuvs::core::from_dlpack(upper_intervals_tensor); + auto corrections = cuvs::core::from_dlpack(additional_corrections_tensor); + auto sums = cuvs::core::from_dlpack(quantized_component_sums_tensor); + auto centroid = cuvs::core::from_dlpack(centroid_tensor); + auto delta = cuvs::core::from_dlpack(dequant_delta_tensor); + auto sum_delta = cuvs::core::from_dlpack(dequant_sum_delta_tensor); + auto row_norm = cuvs::core::from_dlpack(row_norm_tensor); + auto layout = to_cpp_layout(c_layout); + + const auto n_rows = codes.extent(0); + const auto dim = static_cast(centroid.extent(0)); + RAFT_EXPECTS(n_rows > 0 && dim > 0, "cuvsBbqQuantizerCreateView: quantizer must not be empty"); + RAFT_EXPECTS( + codes.extent(1) == + static_cast( + cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dim, layout)), + "cuvsBbqQuantizerCreateView: code row length does not match the dimension and layout"); + RAFT_EXPECTS(lower.extent(0) == n_rows && upper.extent(0) == n_rows && + corrections.extent(0) == n_rows && sums.extent(0) == n_rows && + delta.extent(0) == n_rows && sum_delta.extent(0) == n_rows && + row_norm.extent(0) == n_rows, + "cuvsBbqQuantizerCreateView: every correction vector must contain one value per row"); + + return quantizer_view_t{codes, + lower, + upper, + corrections, + sums, + centroid, + delta, + sum_delta, + row_norm, + layout, + static_cast(metric), + centroid_norm_sq}; +} + +template +void make_and_bind_quantizer_view( + cuvsBbqQuantizer_t* output, + DLDataType dtype, + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq) +{ + using view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + auto view = std::make_unique(make_cpp_quantizer_view(codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq)); + auto handle = std::make_unique(); + handle->addr = reinterpret_cast(view.release()); + handle->destroy_addr = &destroy_typed_addr; + handle->dtype = dtype; + handle->is_owning = false; + *output = handle.release(); +} + +} // namespace + +extern "C" cuvsError_t cuvsBbqQuantizerCreateView( + DLManagedTensor* codes, + DLManagedTensor* lower_intervals, + DLManagedTensor* upper_intervals, + DLManagedTensor* additional_corrections, + DLManagedTensor* quantized_component_sums, + DLManagedTensor* centroid, + DLManagedTensor* dequant_delta, + DLManagedTensor* dequant_sum_delta, + DLManagedTensor* row_norm, + cuvsBbqCodeLayout_t layout, + cuvsDistanceType metric, + float centroid_norm_sq, + cuvsBbqQuantizer_t* quantizer) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(quantizer != nullptr, "cuvsBbqQuantizerCreateView: null output"); + *quantizer = nullptr; + RAFT_EXPECTS(codes != nullptr && lower_intervals != nullptr && upper_intervals != nullptr && + additional_corrections != nullptr && quantized_component_sums != nullptr && + centroid != nullptr && dequant_delta != nullptr && + dequant_sum_delta != nullptr && row_norm != nullptr, + "cuvsBbqQuantizerCreateView: null tensor"); + auto dtype = centroid->dl_tensor.dtype; + if (dtype.code == kDLFloat && dtype.bits == 32) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLFloat && dtype.bits == 16) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLInt && dtype.bits == 8) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + make_and_bind_quantizer_view(quantizer, + dtype, + codes, + lower_intervals, + upper_intervals, + additional_corrections, + quantized_component_sums, + centroid, + dequant_delta, + dequant_sum_delta, + row_norm, + layout, + metric, + centroid_norm_sq); + } else { + RAFT_FAIL("cuvsBbqQuantizerCreateView: unsupported centroid dtype: code=%d, bits=%d", + dtype.code, + dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer) +{ + return cuvs::core::translate_exceptions([=] { + if (quantizer == nullptr) { return; } + if (quantizer->destroy_addr != nullptr && quantizer->addr != 0) { + quantizer->destroy_addr(reinterpret_cast(quantizer->addr)); + } + delete quantizer; + }); +} diff --git a/c/tests/core/dataset_c.cu b/c/tests/core/dataset_c.cu index 0d05c497ff..bebf3f3196 100644 --- a/c/tests/core/dataset_c.cu +++ b/c/tests/core/dataset_c.cu @@ -3,13 +3,17 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/core/interop.hpp" + #include #include #include +#include #include #include #include +#include #include #include @@ -37,6 +41,27 @@ struct MatrixTensor { } }; +template +auto make_device_matrix_tensor(T* data, int64_t rows, int64_t columns) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_matrix_view(data, rows, columns), &tensor); + return tensor; +} + +template +auto make_device_vector_tensor(T* data, int64_t size) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& tensor) +{ + if (tensor.deleter != nullptr) { tensor.deleter(&tensor); } +} + } // namespace TEST(DatasetC, CreateDestroy) @@ -213,3 +238,109 @@ TEST(DatasetC, MakeHostPaddedFromDevice) ASSERT_EQ(cuvsDatasetDestroy(padded), CUVS_SUCCESS); ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); } + +TEST(DatasetC, MakeBbqView) +{ + constexpr int64_t n_rows = 4; + constexpr int64_t dim = 8; + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cudaStream_t stream; + ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); + + rmm::device_uvector codes_1b(n_rows, stream); + rmm::device_uvector lower(n_rows, stream); + rmm::device_uvector upper(n_rows, stream); + rmm::device_uvector corrections(n_rows, stream); + rmm::device_uvector sums(n_rows, stream); + rmm::device_uvector centroid(dim, stream); + rmm::device_uvector delta(n_rows, stream); + rmm::device_uvector sum_delta(n_rows, stream); + rmm::device_uvector row_norm(n_rows, stream); + + auto codes_tensor = make_device_matrix_tensor(codes_1b.data(), n_rows, 1); + auto lower_tensor = make_device_vector_tensor(lower.data(), n_rows); + auto upper_tensor = make_device_vector_tensor(upper.data(), n_rows); + auto corrections_tensor = make_device_vector_tensor(corrections.data(), n_rows); + auto sums_tensor = make_device_vector_tensor(sums.data(), n_rows); + auto centroid_tensor = make_device_vector_tensor(centroid.data(), dim); + auto delta_tensor = make_device_vector_tensor(delta.data(), n_rows); + auto sum_delta_tensor = make_device_vector_tensor(sum_delta.data(), n_rows); + auto row_norm_tensor = make_device_vector_tensor(row_norm.data(), n_rows); + + cuvsBbqQuantizer_t quantizer; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_PACKED_1B, + L2Expanded, + 0.0f, + &quantizer), + CUVS_SUCCESS); + + cuvsDataset_t dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, &quantizer, 1, &dataset), CUVS_SUCCESS); + + cuvsDatasetLayout_t layout; + cuvsDatasetMemType_t mem_type; + bool is_owning; + DLDataType dtype; + ASSERT_EQ(cuvsDatasetGetLayout(dataset, &layout), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetMemType(dataset, &mem_type), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetIsOwning(dataset, &is_owning), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetDtype(dataset, &dtype), CUVS_SUCCESS); + EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_BBQ); + EXPECT_EQ(mem_type, CUVS_DATASET_MEM_TYPE_DEVICE); + EXPECT_FALSE(is_owning); + EXPECT_EQ(dtype.code, kDLFloat); + EXPECT_EQ(dtype.bits, 32); + + ASSERT_EQ(cuvsDatasetDestroy(dataset), CUVS_SUCCESS); + + rmm::device_uvector codes_4t(n_rows * 4, stream); + auto codes_4t_tensor = make_device_matrix_tensor(codes_4t.data(), n_rows, 4); + cuvsBbqQuantizer_t quantizer_4t; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_4t_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + L2Expanded, + 0.0f, + &quantizer_4t), + CUVS_SUCCESS); + + cuvsBbqQuantizer_t quantizers[] = {quantizer, quantizer_4t}; + cuvsDataset_t asymmetric_dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, quantizers, 2, &asymmetric_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetGetLayout(asymmetric_dataset, &layout), CUVS_SUCCESS); + EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_BBQ); + + ASSERT_EQ(cuvsDatasetDestroy(asymmetric_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer_4t), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); + + free_tensor(codes_tensor); + free_tensor(codes_4t_tensor); + free_tensor(lower_tensor); + free_tensor(upper_tensor); + free_tensor(corrections_tensor); + free_tensor(sums_tensor); + free_tensor(centroid_tensor); + free_tensor(delta_tensor); + free_tensor(sum_delta_tensor); + free_tensor(row_norm_tensor); +} diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index ba0e3ee310..e560962a42 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -4,10 +4,13 @@ */ #include "test_utils.cuh" +#include "../../src/core/interop.hpp" #include #include #include +#include #include +#include #include #include @@ -36,6 +39,31 @@ #include +namespace { + +template +auto make_device_matrix_tensor(T* data, int64_t rows, int64_t columns) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_matrix_view(data, rows, columns), &tensor); + return tensor; +} + +template +auto make_device_vector_tensor(T* data, int64_t size) -> DLManagedTensor +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& tensor) +{ + if (tensor.deleter != nullptr) { tensor.deleter(&tensor); } +} + +} // namespace + float dataset[4][2] = {{0.74021935, 0.9209938}, {0.03902049, 0.9689629}, {0.92514056, 0.4463501}, @@ -179,6 +207,153 @@ TEST(CagraC, BuildSearch) cuvsResourcesDestroy(res); } +TEST(CagraC, BuildBbqGraphAndAttachDataset) +{ + constexpr int64_t n_rows = 64; + constexpr int64_t dim = 32; + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cudaStream_t stream; + ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); + + rmm::device_uvector codes(n_rows * dim, stream); + rmm::device_uvector lower(n_rows, stream); + rmm::device_uvector upper(n_rows, stream); + rmm::device_uvector corrections(n_rows, stream); + rmm::device_uvector sums(n_rows, stream); + rmm::device_uvector centroid(dim, stream); + rmm::device_uvector delta(n_rows, stream); + rmm::device_uvector sum_delta(n_rows, stream); + rmm::device_uvector row_norm(n_rows, stream); + rmm::device_uvector dense(n_rows * dim, stream); + ASSERT_EQ(cudaMemsetAsync(codes.data(), 0, codes.size(), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(lower.data(), 0, lower.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(upper.data(), 0, upper.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(corrections.data(), 0, corrections.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(sums.data(), 0, sums.size() * sizeof(int32_t), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(centroid.data(), 0, centroid.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(delta.data(), 0, delta.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(sum_delta.data(), 0, sum_delta.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ( + cudaMemsetAsync(row_norm.data(), 0, row_norm.size() * sizeof(float), stream), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(dense.data(), 0, dense.size() * sizeof(float), stream), cudaSuccess); + + auto codes_tensor = make_device_matrix_tensor(codes.data(), n_rows, dim); + auto lower_tensor = make_device_vector_tensor(lower.data(), n_rows); + auto upper_tensor = make_device_vector_tensor(upper.data(), n_rows); + auto corrections_tensor = make_device_vector_tensor(corrections.data(), n_rows); + auto sums_tensor = make_device_vector_tensor(sums.data(), n_rows); + auto centroid_tensor = make_device_vector_tensor(centroid.data(), dim); + auto delta_tensor = make_device_vector_tensor(delta.data(), n_rows); + auto sum_delta_tensor = make_device_vector_tensor(sum_delta.data(), n_rows); + auto row_norm_tensor = make_device_vector_tensor(row_norm.data(), n_rows); + auto dense_tensor = make_device_matrix_tensor(dense.data(), n_rows, dim); + + cuvsBbqQuantizer_t quantizer; + ASSERT_EQ(cuvsBbqQuantizerCreateView(&codes_tensor, + &lower_tensor, + &upper_tensor, + &corrections_tensor, + &sums_tensor, + ¢roid_tensor, + &delta_tensor, + &sum_delta_tensor, + &row_norm_tensor, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B, + L2Expanded, + 0.0f, + &quantizer), + CUVS_SUCCESS); + + cuvsDataset_t bbq_dataset; + ASSERT_EQ(cuvsDatasetMakeBbqView(res, &quantizer, 1, &bbq_dataset), CUVS_SUCCESS); + + cuvsCagraIndexParams_t params; + ASSERT_EQ(cuvsCagraIndexParamsCreate(¶ms), CUVS_SUCCESS); + params->metric = L2Expanded; + params->build_algo = NN_DESCENT; + params->graph_degree = 8; + params->intermediate_graph_degree = 16; + params->nn_descent_niter = 5; + + cuvsCagraIndex_t index; + ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraBuild(res, params, bbq_dataset, index), CUVS_SUCCESS); + + DLManagedTensor graph_tensor{}; + ASSERT_EQ(cuvsCagraIndexGetGraph(index, &graph_tensor), CUVS_SUCCESS); + ASSERT_EQ(graph_tensor.dl_tensor.shape[0], n_rows); + ASSERT_EQ(graph_tensor.dl_tensor.shape[1], params->graph_degree); + std::vector graph(n_rows * params->graph_degree); + raft::copy(graph.data(), + static_cast(graph_tensor.dl_tensor.data), + graph.size(), + stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + for (auto neighbor : graph) { + EXPECT_LT(neighbor, static_cast(n_rows)); + } + + rmm::device_uvector query(dim, stream); + rmm::device_uvector neighbors(1, stream); + rmm::device_uvector distances(1, stream); + ASSERT_EQ(cudaMemsetAsync(query.data(), 0, query.size() * sizeof(float), stream), cudaSuccess); + auto query_tensor = make_device_matrix_tensor(query.data(), 1, dim); + auto neighbors_tensor = make_device_matrix_tensor(neighbors.data(), 1, 1); + auto distances_tensor = make_device_matrix_tensor(distances.data(), 1, 1); + cuvsCagraSearchParams_t search_params; + ASSERT_EQ(cuvsCagraSearchParamsCreate(&search_params), CUVS_SUCCESS); + cuvsFilter filter{.addr = 0, .type = NO_FILTER}; + EXPECT_EQ(cuvsCagraSearch(res, + search_params, + index, + &query_tensor, + &neighbors_tensor, + &distances_tensor, + filter), + CUVS_ERROR); + EXPECT_NE(std::string(cuvsGetLastErrorText()).find("BBQ-built"), std::string::npos); + + cuvsDataset_t padded_dataset; + ASSERT_EQ(cuvsDatasetMakePaddedView(res, &dense_tensor, &padded_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraUpdateDataset(res, padded_dataset, index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraSearch(res, + search_params, + index, + &query_tensor, + &neighbors_tensor, + &distances_tensor, + filter), + CUVS_SUCCESS); + + ASSERT_EQ(cuvsCagraSearchParamsDestroy(search_params), CUVS_SUCCESS); + if (graph_tensor.deleter != nullptr) { graph_tensor.deleter(&graph_tensor); } + ASSERT_EQ(cuvsCagraIndexDestroy(index), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetDestroy(padded_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraIndexParamsDestroy(params), CUVS_SUCCESS); + ASSERT_EQ(cuvsDatasetDestroy(bbq_dataset), CUVS_SUCCESS); + ASSERT_EQ(cuvsBbqQuantizerDestroy(quantizer), CUVS_SUCCESS); + ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); + + free_tensor(codes_tensor); + free_tensor(lower_tensor); + free_tensor(upper_tensor); + free_tensor(corrections_tensor); + free_tensor(sums_tensor); + free_tensor(centroid_tensor); + free_tensor(delta_tensor); + free_tensor(sum_delta_tensor); + free_tensor(row_norm_tensor); + free_tensor(dense_tensor); + free_tensor(query_tensor); + free_tensor(neighbors_tensor); + free_tensor(distances_tensor); +} + // CAGRA operations that need a search-ready index must reject host / non-device-padded // datasets rather than succeeding and producing undefined behavior. TEST(CagraC, DatasetContractFailures) diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 8ed1b8cd8b..30de27cfc8 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -21,11 +21,11 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcuvs" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '365Mi' + --max-allowed-size-compressed '375Mi' ) else PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '225Mi' + --max-allowed-size-compressed '235Mi' ) fi elif [[ "${package_dir}" != "python/cuvs" ]]; then diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 1beaf2484f..ef9977f72c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1533,6 +1533,7 @@ if(NOT BUILD_CPU_ONLY) src/preprocessing/quantize/scalar.cu src/preprocessing/quantize/binary.cu src/preprocessing/quantize/pq.cu + src/preprocessing/quantize/bbq.cu src/preprocessing/spectral/spectral_embedding.cu src/preprocessing/pca/pca.cu ${select_k_inst_files} @@ -1838,9 +1839,10 @@ endif() # ################################################################################################## # * build test executable ---------------------------------------------------- +add_subdirectory(internal) + if(BUILD_TESTS) enable_testing() - add_subdirectory(internal) add_subdirectory(tests) endif() diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 204c0d3ffb..9969d6bf95 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -956,6 +957,10 @@ template using device_pq_index = index>; +/** CAGRA index with a device-resident BBQ-quantized dataset. */ +template +using device_bbq_index = index>; + /** Index type returned by `cagra::build(res, params, dataset_view)`. */ template using cagra_index_t = @@ -1214,6 +1219,49 @@ auto build(raft::resources const& res, cuvs::neighbors::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; +/** + * @brief Build from a device BBQ-quantized dataset view. + * + * The kNN graph is built from the quantized codes alone, so the uncompressed vectors are never + * needed and peak memory is driven by the code size. Only nn-descent graph construction is + * available (IVF-PQ, iterative CAGRA search, and ACE all read uncompressed vectors), and the + * metric must be one of L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct and must match + * the metric the quantizer corrections were generated for. + * + * The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the type-changing + * `update_dataset` with an uncompressed device-padded dataset to search the resulting graph. + * + * @param[in] res raft resources + * @param[in] params CAGRA index build parameters + * @param[in] dataset device BBQ dataset view [n_rows, dim] + * @return built `device_bbq_index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + +/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, + * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::device_bbq_index; + /** * @} */ @@ -2965,6 +3013,75 @@ void deserialize(raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_pq_index* index); +/* device_bbq_index graph-only overloads (CAGRA + BBQ). + * + * These overloads persist the graph and index metadata, but not the attached BBQ dataset + * Attach a compatible dataset with `update_dataset` after deserialization before searching. + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_bbq_index* index); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::device_bbq_index& index); + +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_bbq_index* index); + /** @copydoc serialize */ void serialize(raft::resources const& handle, const std::string& filename, @@ -4948,6 +5065,48 @@ auto update_dataset(raft::resources const& res, device_vpq_dataset_view dataset) -> index>; +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + device_padded_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset(raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + device_vpq_dataset_view dataset) + -> index>; + } // namespace cagra } // namespace neighbors } // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 0db1485a2f..baa5b44da2 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -779,6 +779,7 @@ enum class dataset_view_kind { standard, vpq_f16, vpq_f32, + bbq, }; /** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 4c031049e2..1e1546c651 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,11 +1,12 @@ /* - * 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 */ #pragma once #include +#include #include #include @@ -523,6 +524,58 @@ auto build(raft::resources const& res, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @brief Build an NN-Descent index directly from a device-resident BBQ dataset. + * + * Symmetric compressed-code distances are used during graph construction. Supported metrics are + * L2Expanded, L2SqrtExpanded, CosineExpanded, and InnerProduct. The dataset's correction terms + * must have been generated for the selected metric. + * + * @param res raft resources + * @param params NN-Descent build parameters + * @param dataset BBQ codes and correction terms in device memory + * @param graph optional caller-owned host graph + * @return index containing the all-neighbors graph + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** + * @copydoc build(raft::resources const& res, index_params const& params, + * cuvs::neighbors::device_bbq_dataset_view dataset, + * std::optional> graph = std::nullopt) + * -> cuvs::neighbors::nn_descent::index + */ +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph = + std::nullopt) -> cuvs::neighbors::nn_descent::index; + /** @} */ /** * @brief Test if we have enough GPU memory to run NN descent algorithm. diff --git a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp new file mode 100644 index 0000000000..21366d47d0 --- /dev/null +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -0,0 +1,388 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +namespace preprocessing::quantize::bbq { + +/** + * @defgroup bbq Better Binary Quantization utilities + * @{ + */ + +/** + * Storage layout of BBQ/OSQ quantized component codes in each dataset row. + * packed_1b: Each dimension is quantized to a single bit and packed into bytes. Reflects + * Lucene's OptimizedScalarQuantizer.packAsBinary. + * transposed_2b: Each dimension is quantized to 2 bits, stored as 2 bitplanes. + * Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only + * (paired with a transposed_4b or packed_1b operand); + * transposed_4b: Each dimension is quantized to 4 bits, optimized for bitwise operations. + * Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of + * every dimension is in the first set dimensions bits, or (dimensions/8) + * bytes. The second, third, and fourth bits are in the second, third, and + * fourth set of dimensions bits, respectively. Format used for queries. + * packed_4b: Each dimension is quantized to 4 bits, two values are packed into each output + * byte. + * packed_7b: Each dimension is quantized to 7 bits and treated as a signed value. + * packed_8b: Each dimension is quantized to 8 bits and treated as an unsigned value. + */ +enum class bbq_code_layout { + packed_1b, + transposed_2b, + transposed_4b, + packed_4b, + packed_7b, + packed_8b, +}; + +/** + * Bit width of a layout. + */ +constexpr auto get_bit_width(bbq_code_layout layout) noexcept -> uint32_t +{ + switch (layout) { + case bbq_code_layout::packed_1b: return 1; + case bbq_code_layout::transposed_2b: return 2; + case bbq_code_layout::transposed_4b: + case bbq_code_layout::packed_4b: return 4; + case bbq_code_layout::packed_7b: return 7; + case bbq_code_layout::packed_8b: return 8; + } + return 0; +} + +/** Bytes one row of @p dim components occupies once encoded in @p layout. */ +constexpr auto get_encoded_row_length(uint32_t dim, bbq_code_layout layout) noexcept -> uint32_t +{ + switch (layout) { + case bbq_code_layout::packed_1b: return raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::transposed_2b: return 2 * raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::transposed_4b: return 4 * raft::div_rounding_up_safe(dim, 8u); + case bbq_code_layout::packed_4b: return raft::div_rounding_up_safe(dim, 2u); + case bbq_code_layout::packed_7b: return dim; + case bbq_code_layout::packed_8b: return dim; + } + return 0; +} + +template +struct quantizer_view; + +/** + * @brief Better Binary Quantization + * ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) + * is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas + * introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et + * al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective + * factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements + * this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch + * exposes it as BBQ. + * + * BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared + * centroid, no random rotation, and OSQ codes. + * + * RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain + * additional per-vector information so search is better than naïve sign-bit comparison. They differ + * in transformation and scale representation. RaBitQ commonly separates residual magnitude from + * direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector + * scalar intervals to interpret the compressed residual codes. + */ +template +struct quantizer { + raft::device_matrix codes; + raft::device_vector lower_intervals; + raft::device_vector upper_intervals; + raft::device_vector additional_corrections; + raft::device_vector quantized_component_sums; + raft::device_vector centroid; + /** Precomputed per-row dequantization factors, derived once (offline) from + * lower/upper_intervals and quantized_component_sums: dequant_delta = (upper-lower)/(2^bits-1) */ + raft::device_vector dequant_delta; + /** Precomputed per-row dequantization factors, derived once (offline) from dequant_delta and + * quantized_component_sums: dequant_sum_delta = dequant_delta * quantized_component_sums. */ + raft::device_vector dequant_sum_delta; + /** Squared norm of the row in original (un-centered) vector space, ||x||^2 */ + raft::device_vector row_norm; + + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + quantizer(raft::resources const& res, + IdxT n_rows, + uint32_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) + : codes{raft::make_device_matrix( + res, n_rows, static_cast(get_encoded_row_length(dim, layout)))}, + lower_intervals{raft::make_device_vector(res, n_rows)}, + upper_intervals{raft::make_device_vector(res, n_rows)}, + additional_corrections{raft::make_device_vector(res, n_rows)}, + quantized_component_sums{raft::make_device_vector(res, n_rows)}, + centroid{raft::make_device_vector(res, static_cast(dim))}, + dequant_delta{raft::make_device_vector(res, n_rows)}, + dequant_sum_delta{raft::make_device_vector(res, n_rows)}, + row_norm{raft::make_device_vector(res, n_rows)}, + layout{layout}, + metric{metric} + { + } + + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return codes.extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(centroid.extent(0)); + } + [[nodiscard]] constexpr auto encoded_row_length() const noexcept -> uint32_t + { + return get_encoded_row_length(dim(), layout); + } + [[nodiscard]] auto view() const noexcept -> quantizer_view + { + return quantizer_view{codes.view(), + lower_intervals.view(), + upper_intervals.view(), + additional_corrections.view(), + quantized_component_sums.view(), + centroid.view(), + dequant_delta.view(), + dequant_sum_delta.view(), + row_norm.view(), + layout, + metric, + centroid_norm_sq}; + } +}; + +template +struct quantizer_view { + raft::device_matrix_view codes; + raft::device_vector_view lower_intervals; + raft::device_vector_view upper_intervals; + raft::device_vector_view additional_corrections; + raft::device_vector_view quantized_component_sums; + raft::device_vector_view centroid; + raft::device_vector_view dequant_delta; + raft::device_vector_view dequant_sum_delta; + raft::device_vector_view row_norm; + + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; + + quantizer_view(raft::device_matrix_view codes_, + raft::device_vector_view lower_intervals_, + raft::device_vector_view upper_intervals_, + raft::device_vector_view additional_corrections_, + raft::device_vector_view quantized_component_sums_, + raft::device_vector_view centroid_, + raft::device_vector_view dequant_delta_, + raft::device_vector_view dequant_sum_delta_, + raft::device_vector_view row_norm_, + bbq_code_layout layout_, + cuvs::distance::DistanceType metric_, + float centroid_norm_sq_) noexcept + : codes{codes_}, + lower_intervals{lower_intervals_}, + upper_intervals{upper_intervals_}, + additional_corrections{additional_corrections_}, + quantized_component_sums{quantized_component_sums_}, + centroid{centroid_}, + dequant_delta{dequant_delta_}, + dequant_sum_delta{dequant_sum_delta_}, + row_norm{row_norm_}, + layout{layout_}, + metric{metric_}, + centroid_norm_sq{centroid_norm_sq_} + { + } + + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT { return codes.extent(0); } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return static_cast(centroid.extent(0)); + } +}; + +namespace helpers { +/** + * Derives dequant_delta from lower/upper_intervals and the layout's code width, and + * dequant_sum_delta from that delta and quantized_component_sums. + */ +void resolve_dequant_factors( + raft::resources const& res, + raft::device_vector_view dequant_delta, + raft::device_vector_view dequant_sum_delta, + raft::device_vector_view lower_intervals, + raft::device_vector_view upper_intervals, + raft::device_vector_view quantized_component_sums, + bbq_code_layout layout); +} // namespace helpers +/** @} */ // end of bbq group + +} // namespace preprocessing::quantize::bbq + +namespace neighbors { +struct bbq_dataset_container { + template + using owning_storage = cuvs::preprocessing::quantize::bbq::quantizer; + template + using view_storage = cuvs::preprocessing::quantize::bbq::quantizer_view; +}; + +template +struct dataset { + using owning_storage_type = bbq_dataset_container::owning_storage; + std::vector quantizers; + + dataset(owning_storage_type&& quantizer) noexcept { add_quantizer(std::move(quantizer)); } + [[nodiscard]] auto as_dataset_view() const noexcept + -> dataset_view> + { + return dataset_view>{quantizers}; + } + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT + { + return quantizers.size() > 0 ? quantizers[0].n_rows() : 0; + } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return quantizers.size() > 0 ? quantizers[0].dim() : 0; + } + + void add_quantizer(owning_storage_type&& quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(std::move(quantizer)); + } + bool has_layout(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return true; } + } + return false; + } +}; + +template +struct dataset_view { + using owning_storage_type = bbq_dataset_container::owning_storage; + using view_storage_type = bbq_dataset_container::view_storage; + std::vector quantizers; + + dataset_view() noexcept = default; + + dataset_view(const std::vector& quantizers) noexcept + { + for (const auto& quantizer : quantizers) { + add_quantizer(quantizer); + } + } + [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT + { + return quantizers.size() > 0 ? quantizers[0].n_rows() : 0; + } + [[nodiscard]] constexpr auto dim() const noexcept -> uint32_t + { + return quantizers.size() > 0 ? quantizers[0].dim() : 0; + } + + void add_quantizer(view_storage_type quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(quantizer); + } + void add_quantizer(const owning_storage_type& quantizer) + { + RAFT_EXPECTS(!has_layout(quantizer.layout), "Quantizer already exists with layout."); + quantizers.push_back(quantizer.view()); + } + bool has_layout(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const noexcept + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return true; } + } + return false; + } + view_storage_type get_quantizer(cuvs::preprocessing::quantize::bbq::bbq_code_layout layout) const + { + for (uint32_t i = 0; i < quantizers.size(); i++) { + if (quantizers[i].layout == layout) { return quantizers[i]; } + } + throw std::runtime_error("No quantizer found with layout."); + } +}; + +template +using device_bbq_dataset = + dataset>; + +template +using device_bbq_dataset_view = + dataset_view>; + +template +struct owning_dataset_for_view> { + using type = device_bbq_dataset; +}; + +template +struct is_bbq_dataset : std::false_type {}; + +template +struct is_bbq_dataset> : std::true_type {}; + +template +inline constexpr bool is_bbq_dataset_v = is_bbq_dataset::value; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = dataset_view_kind::bbq; +}; + +template +struct cagra_view_element_type> { + using type = DataT; +}; + +template +inline constexpr bool is_device_bbq_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::bbq && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_bbq_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::bbq && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_bbq_dataset_view_v = + is_device_bbq_dataset_view_v || is_host_bbq_dataset_view_v; + +} // namespace neighbors + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/internal/CMakeLists.txt b/cpp/internal/CMakeLists.txt index 4a50bfef58..403d06310d 100644 --- a/cpp/internal/CMakeLists.txt +++ b/cpp/internal/CMakeLists.txt @@ -1,14 +1,17 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= -if(BUILD_TESTS OR BUILD_MICRO_BENCH) +if(BUILD_TESTS + OR BUILD_MICRO_BENCH + OR BUILD_CUVS_BENCH +) add_library(cuvs_internal INTERFACE) target_include_directories( - cuvs_internal INTERFACE "$" + cuvs_internal INTERFACE "$" ) target_compile_features(cuvs_internal INTERFACE cxx_std_17 $) endif() diff --git a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp new file mode 100644 index 0000000000..603517e518 --- /dev/null +++ b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp @@ -0,0 +1,629 @@ +/* + * SPDX-FileCopyrightText: Copyright the Apache Software Foundation (ASF) + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * Host-side reference implementation of Lucene's OptimizedScalarQuantizer, used to feed the + * BBQ NN-Descent build path. This is prototype scaffolding: it is a slow CPU quantizer that + * exists so the compressed build can be exercised end to end, and it is expected to be replaced + * by a GPU quantizer. It is shared between the unit tests and the ann-bench CAGRA wrapper so the + * two can never disagree about the code format; libcuvs itself only ever sees the uploaded codes. + */ +namespace cuvs_internal::bbq { + +using cuvs::preprocessing::quantize::bbq::bbq_code_layout; +using cuvs::preprocessing::quantize::bbq::get_bit_width; +using cuvs::preprocessing::quantize::bbq::get_encoded_row_length; + +/** + * Host-resident mirror of the arrays in a BBQ quantizer. The library quantizer is device-only, + * so the CPU reference implementation and its on-disk cache need their own staging type; it is + * uploaded by copy_bbq_owning_storage_host_to_device below. + */ +struct host_quantizer_storage { + raft::host_matrix codes; + raft::host_vector lower_intervals; + raft::host_vector upper_intervals; + raft::host_vector additional_corrections; + raft::host_vector quantized_component_sums; + raft::host_vector centroid; + raft::host_vector dequant_delta; + raft::host_vector dequant_sum_delta; + raft::host_vector row_norm; + bbq_code_layout layout{bbq_code_layout::packed_1b}; + cuvs::distance::DistanceType metric{cuvs::distance::DistanceType::L2Expanded}; + float centroid_norm_sq{}; +}; + +/** + * Fills the derived per-row dequantization factors. The device quantizer requires them up front + * (deriving them there would need a stream to allocate with), so they are computed here, where + * the intervals are already host-resident. + */ +inline void derive_dequant_factors(host_quantizer_storage& q) +{ + const float scale = 1.0f / static_cast((uint32_t{1} << get_bit_width(q.layout)) - 1); + for (int64_t i = 0; i < q.lower_intervals.extent(0); ++i) { + const float delta = (q.upper_intervals(i) - q.lower_intervals(i)) * scale; + q.dequant_delta(i) = delta; + q.dequant_sum_delta(i) = delta * static_cast(q.quantized_component_sums(i)); + } +} + +constexpr float kMinimumMseGrid[8][2] = {{-0.798f, 0.798f}, + {-1.493f, 1.493f}, + {-2.051f, 2.051f}, + {-2.514f, 2.514f}, + {-2.916f, 2.916f}, + {-3.278f, 3.278f}, + {-3.611f, 3.611f}, + {-3.922f, 3.922f}}; + +constexpr float kDefaultLambda = 0.1f; +constexpr int kDefaultIters = 5; + +inline long round(double x) { return static_cast(std::floor(x + 0.5)); } + +inline double clamp(double x, double a, double b) { return std::min(std::max(x, a), b); } + +inline double loss( + const std::vector& vector, const float interval[2], int points, float norm2, float lambda) +{ + const double a = interval[0]; + const double b = interval[1]; + const double step = (b - a) / (points - 1.0); + const double step_inv = 1.0 / step; + double xe = 0.0; + double e = 0.0; + for (double xi : vector) { + const double xiq = a + step * static_cast(round((clamp(xi, a, b) - a) * step_inv)); + xe += xi * (xi - xiq); + e += (xi - xiq) * (xi - xiq); + } + return (1.0 - lambda) * xe * xe / norm2 + lambda * e; +} + +inline void optimize_intervals(float interval[2], + const std::vector& vector, + float norm2, + int points, + float lambda = kDefaultLambda, + int iters = kDefaultIters) +{ + double initial_loss = loss(vector, interval, points, norm2, lambda); + const float scale = (1.0f - lambda) / norm2; + if (!std::isfinite(scale)) { return; } + for (int i = 0; i < iters; ++i) { + const float a = interval[0]; + const float b = interval[1]; + const float step_inv = (points - 1.0f) / (b - a); + double daa = 0.0, dab = 0.0, dbb = 0.0, dax = 0.0, dbx = 0.0; + for (float xi : vector) { + const float k = + static_cast(round(static_cast((clamp(xi, a, b) - a) * step_inv))); + const float s = k / (points - 1); + daa += (1.0 - s) * (1.0 - s); + dab += (1.0 - s) * s; + dbb += s * s; + dax += xi * (1.0 - s); + dbx += xi * s; + } + const double m0 = scale * dax * dax + lambda * daa; + const double m1 = scale * dax * dbx + lambda * dab; + const double m2 = scale * dbx * dbx + lambda * dbb; + const double det = m0 * m2 - m1 * m1; + if (det == 0) { return; } + const float a_opt = static_cast((m2 * dax - m1 * dbx) / det); + const float b_opt = static_cast((m0 * dbx - m1 * dax) / det); + if (std::abs(interval[0] - a_opt) < 1e-8 && std::abs(interval[1] - b_opt) < 1e-8) { return; } + float new_interval[2] = {a_opt, b_opt}; + const double new_loss = loss(vector, new_interval, points, norm2, lambda); + if (new_loss > initial_loss) { return; } + interval[0] = a_opt; + interval[1] = b_opt; + initial_loss = new_loss; + } +} + +struct row_result { + float lower_interval; + float upper_interval; + float additional_correction; + int32_t quantized_component_sum; +}; + +inline row_result scalar_quantize(std::vector& vector, + std::vector& destination, + uint8_t bits, + const float* centroid, + bool euclidean) +{ + const int n = static_cast(vector.size()); + const int points = 1 << bits; + double vec_mean = 0.0; + double vec_var = 0.0; + float norm2 = 0.0f; + float centroid_dot = 0.0f; + float min = FLT_MAX; + float max = -FLT_MAX; + for (int i = 0; i < n; ++i) { + if (!euclidean) { centroid_dot += vector[i] * centroid[i]; } + vector[i] = vector[i] - centroid[i]; + min = std::min(min, vector[i]); + max = std::max(max, vector[i]); + norm2 += vector[i] * vector[i]; + const double delta = vector[i] - vec_mean; + vec_mean += delta / (i + 1); + vec_var += delta * (vector[i] - vec_mean); + } + vec_var /= n; + const double vec_std = std::sqrt(vec_var); + + float interval[2]; + interval[0] = + static_cast(clamp(kMinimumMseGrid[bits - 1][0] * vec_std + vec_mean, min, max)); + interval[1] = + static_cast(clamp(kMinimumMseGrid[bits - 1][1] * vec_std + vec_mean, min, max)); + optimize_intervals(interval, vector, norm2, points); + + const float n_steps = static_cast((1 << bits) - 1); + const float a = interval[0]; + const float b = interval[1]; + const float step = (b - a) / n_steps; + int sum_query = 0; + for (int h = 0; h < n; ++h) { + const float xi = static_cast(clamp(vector[h], a, b)); + const int assignment = static_cast(round((xi - a) / step)); + sum_query += assignment; + destination[h] = static_cast(assignment); + } + return row_result{interval[0], interval[1], euclidean ? norm2 : centroid_dot, sum_query}; +} + +// Packs one-byte-per-component codes into packed_1b / transposed_2b / +// packed_4b / transposed_4b (or leaves unpacked). Matches Lucene packAsBinary, +// packNibbles, transposeDibit, transposeHalfByte. +inline std::vector pack_codes(const std::vector& unpacked, + size_t n_rows, + size_t dim, + bbq_code_layout layout) +{ + const uint32_t bits = get_bit_width(layout); + const size_t row_length = get_encoded_row_length(static_cast(dim), layout); + if (layout == bbq_code_layout::packed_8b || layout == bbq_code_layout::packed_7b) { + return unpacked; + } + + std::vector packed(n_rows * row_length, 0); +#pragma omp parallel for + for (int64_t row = 0; row < static_cast(n_rows); ++row) { + auto* output = packed.data() + static_cast(row) * row_length; + const auto* input = unpacked.data() + static_cast(row) * dim; + if (layout == bbq_code_layout::packed_4b) { + // Contiguous: dims 2k / 2k+1 share byte k. NOT Lucene packNibbles, which pairs dim i with + // dim dim/2 + i. A self-join is position-agnostic so either works there, but an asymmetric + // pair (packed_1b or transposed_2b document promoted to 4-bit width against this query) + // needs dimension k of both operands in the same slot -- halves-pairing silently + // multiplies mismatched dimensions and costs recall. + const size_t pairs = dim / 2; + for (size_t i = 0; i < pairs; ++i) { + output[i] = static_cast((input[2 * i] << 4) | (input[2 * i + 1] & 0x0f)); + } + continue; + } + for (size_t d = 0; d < dim; ++d) { + const uint8_t code = input[d]; + if (layout == bbq_code_layout::packed_1b) { + for (uint32_t bit = 0; bit < bits; ++bit) { + const size_t position = d * bits + bit; + output[position / 8] |= + static_cast(((code >> (bits - 1 - bit)) & 1u) << (7 - position % 8)); + } + } else { + // transposed_2b / transposed_4b bit-planes (LSB plane first) + const size_t stripe = (dim + 7) / 8; + for (uint32_t bit = 0; bit < bits; ++bit) { + output[bit * stripe + d / 8] |= static_cast(((code >> bit) & 1u) << (7 - d % 8)); + } + } + } + } + return packed; +} + +inline host_quantizer_storage quantize(const float* data, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + bbq_code_layout layout = bbq_code_layout::packed_8b) +{ + const auto bits = static_cast(get_bit_width(layout)); + const bool euclidean = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded; + + auto centroid = raft::make_host_vector(dim); + float centroid_norm_sq = 0.0f; + std::fill_n(centroid.data_handle(), static_cast(dim), 0.0f); + for (int64_t i = 0; i < n_rows; ++i) { + for (int64_t d = 0; d < dim; ++d) { + centroid(d) += data[i * dim + d]; + } + } + for (int64_t d = 0; d < dim; ++d) { + centroid(d) /= static_cast(n_rows); + centroid_norm_sq += centroid(d) * centroid(d); + } + + std::vector unpacked(static_cast(n_rows * dim)); + auto lower_intervals = raft::make_host_vector(n_rows); + auto upper_intervals = raft::make_host_vector(n_rows); + auto additional_corrections = raft::make_host_vector(n_rows); + auto quantized_component_sums = raft::make_host_vector(n_rows); + auto row_norm = raft::make_host_vector(n_rows); + +#pragma omp parallel for + for (int64_t i = 0; i < n_rows; ++i) { + std::vector row(data + i * dim, data + (i + 1) * dim); + std::vector codes(dim); + float orig_norm2 = 0.0f; + for (int64_t d = 0; d < dim; ++d) { + orig_norm2 += row[d] * row[d]; + } + row_norm(i) = orig_norm2; + const auto result = scalar_quantize(row, codes, bits, centroid.data_handle(), euclidean); + std::copy(codes.begin(), codes.end(), unpacked.begin() + i * dim); + lower_intervals(i) = result.lower_interval; + upper_intervals(i) = result.upper_interval; + additional_corrections(i) = result.additional_correction; + quantized_component_sums(i) = result.quantized_component_sum; + } + + auto packed = pack_codes(unpacked, static_cast(n_rows), static_cast(dim), layout); + auto codes = raft::make_host_matrix( + n_rows, get_encoded_row_length(static_cast(dim), layout)); + std::copy(packed.begin(), packed.end(), codes.data_handle()); + + host_quantizer_storage out{std::move(codes), + std::move(lower_intervals), + std::move(upper_intervals), + std::move(additional_corrections), + std::move(quantized_component_sums), + std::move(centroid), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + std::move(row_norm), + layout, + metric, + centroid_norm_sq}; + derive_dequant_factors(out); + return out; +} + +/** The CPU reference quantizer works in float, so it only ever feeds a float-valued dataset. */ +template +auto copy_bbq_owning_storage_host_to_device(raft::resources const& res, + host_quantizer_storage const& host_storage) -> + typename cuvs::neighbors::device_bbq_dataset::owning_storage_type +{ + using device_storage = + typename cuvs::neighbors::device_bbq_dataset::owning_storage_type; + auto stream = raft::resource::get_cuda_stream(res); + device_storage device{res, + static_cast(host_storage.codes.extent(0)), + static_cast(host_storage.centroid.extent(0)), + host_storage.layout, + host_storage.metric}; + + raft::copy( + device.codes.data_handle(), host_storage.codes.data_handle(), device.codes.size(), stream); + raft::copy(device.lower_intervals.data_handle(), + host_storage.lower_intervals.data_handle(), + device.lower_intervals.size(), + stream); + raft::copy(device.upper_intervals.data_handle(), + host_storage.upper_intervals.data_handle(), + device.upper_intervals.size(), + stream); + raft::copy(device.additional_corrections.data_handle(), + host_storage.additional_corrections.data_handle(), + device.additional_corrections.size(), + stream); + raft::copy(device.quantized_component_sums.data_handle(), + host_storage.quantized_component_sums.data_handle(), + device.quantized_component_sums.size(), + stream); + raft::copy(device.centroid.data_handle(), + host_storage.centroid.data_handle(), + device.centroid.size(), + stream); + raft::copy(device.dequant_delta.data_handle(), + host_storage.dequant_delta.data_handle(), + device.dequant_delta.size(), + stream); + raft::copy(device.dequant_sum_delta.data_handle(), + host_storage.dequant_sum_delta.data_handle(), + device.dequant_sum_delta.size(), + stream); + raft::copy(device.row_norm.data_handle(), + host_storage.row_norm.data_handle(), + device.row_norm.size(), + stream); + device.centroid_norm_sq = host_storage.centroid_norm_sq; + return device; +} + +template +auto make_device_bbq_dataset(raft::resources const& res, + std::vector const& host) + -> cuvs::neighbors::device_bbq_dataset +{ + RAFT_EXPECTS(host.size() != 0, "host BBQ dataset has no storage"); + cuvs::neighbors::device_bbq_dataset device{ + copy_bbq_owning_storage_host_to_device(res, host[0])}; + for (std::size_t i = 1; i < host.size(); ++i) { + device.add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); + } + return device; +} + +struct bbq_layout_token { + std::string_view token; + bbq_code_layout layout; +}; + +// Same convention as my_tests/bbq's CLI tokens: bare N = densely packed (tensor-core-eligible), +// N + "t" = transposed/bitplane (SIMT). At 1 bit the two coincide, so there is no "1t". There is +// no densely-packed 2-bit layout at all (packed_2b was retired): transposed_2b is SIMT-only, so +// "2t" is the only 2-bit token and it is not tensor-core-eligible. +constexpr bbq_layout_token kBbqLayoutTokens[] = { + {"1", bbq_code_layout::packed_1b}, + {"2t", bbq_code_layout::transposed_2b}, + {"4", bbq_code_layout::packed_4b}, + {"4t", bbq_code_layout::transposed_4b}, + {"7", bbq_code_layout::packed_7b}, + {"8", bbq_code_layout::packed_8b}, +}; + +/** Parses a layout token. Width follows from the layout (bits=4 is packed_4b vs transposed_4b). */ +inline auto parse_bbq_layout_token(std::string_view token) -> bbq_code_layout +{ + for (const auto& t : kBbqLayoutTokens) { + if (t.token == token) { return t.layout; } + } + RAFT_FAIL("Unknown BBQ layout token '%s'; expected one of 1, 2t, 4, 4t, 7, 8.", + std::string(token).c_str()); +} + +/** + * Reject layout pairs the local-join kernels cannot serve. Without this the library either + * silently falls back to a symmetric join on the first quantizer (which looks like a working + * asymmetric run), or reaches an unsupported-layout failure deep inside GNND::build. Must be kept + * in sync with the dispatch in GNND::local_join (nn_descent.cuh). + */ +inline void validate_layout_pair(bbq_code_layout query_layout, bbq_code_layout doc_layout) +{ + if (query_layout == doc_layout) { + const bool supported = query_layout == bbq_code_layout::packed_1b || + query_layout == bbq_code_layout::transposed_2b || + query_layout == bbq_code_layout::packed_4b || + query_layout == bbq_code_layout::packed_7b || + query_layout == bbq_code_layout::packed_8b; + RAFT_EXPECTS(supported, + "Symmetric BBQ NN-Descent has no local-join kernel for layout %d -- " + "transposed_4b is only supported as one side of an asymmetric pair.", + static_cast(query_layout)); + return; + } + const bool supported = + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::packed_4b) || + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::transposed_2b) || + (doc_layout == bbq_code_layout::packed_1b && query_layout == bbq_code_layout::transposed_4b) || + (doc_layout == bbq_code_layout::transposed_2b && + query_layout == bbq_code_layout::transposed_4b); + RAFT_EXPECTS(supported, + "Asymmetric BBQ NN-Descent supports only (doc, query) layouts of (packed_1b, " + "packed_4b), (packed_1b, transposed_2b), (packed_1b, transposed_4b), or " + "(transposed_2b, transposed_4b); got (%d, %d).", + static_cast(doc_layout), + static_cast(query_layout)); +} + +/** Every parameter that affects the codes is in the name, so the cache self-invalidates. */ +inline auto cache_path(int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> std::string +{ + const char* dir = std::getenv("CUVS_BBQ_CACHE_DIR"); + return std::string{dir != nullptr && dir[0] != '\0' ? dir : "/tmp"} + "/bbq-n" + + std::to_string(n_rows) + "-d" + std::to_string(dim) + "-b" + + std::to_string(get_bit_width(layout)) + "-l" + std::to_string(static_cast(layout)) + + "-m" + std::to_string(static_cast(metric)) + ".bin"; +} + +/** Visits every raw buffer of @p q in a fixed order; this is the on-disk layout. */ +template +void for_each_buffer(host_quantizer_storage& q, OpT op) +{ + op(q.codes.data_handle(), q.codes.size() * sizeof(uint8_t)); + op(q.lower_intervals.data_handle(), q.lower_intervals.size() * sizeof(float)); + op(q.upper_intervals.data_handle(), q.upper_intervals.size() * sizeof(float)); + op(q.additional_corrections.data_handle(), q.additional_corrections.size() * sizeof(float)); + op(q.quantized_component_sums.data_handle(), q.quantized_component_sums.size() * sizeof(int32_t)); + op(q.centroid.data_handle(), q.centroid.size() * sizeof(float)); + // Unlike dequant_delta/dequant_sum_delta, row_norm (original-space ||x||^2) isn't derivable + // from the other cached fields, so it has to round-trip through the cache. + op(q.row_norm.data_handle(), q.row_norm.size() * sizeof(float)); +} + +/** Allocates the arrays of the given shape, leaving their contents undefined. */ +inline auto make_host_quantizer_storage(int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) + -> host_quantizer_storage +{ + return host_quantizer_storage{ + raft::make_host_matrix( + n_rows, get_encoded_row_length(static_cast(dim), layout)), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(dim), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + raft::make_host_vector(n_rows), + layout, + metric, + 0.0f}; +} + +/** Reads the codes back if the file is present and has exactly the expected length. */ +inline auto cache_load(const std::string& path, + int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> std::optional +{ + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) { return std::nullopt; } + auto q = make_host_quantizer_storage(n_rows, dim, layout, metric); + std::streamoff bytes = 0; + for_each_buffer(q, [&bytes](void*, size_t n) { bytes += static_cast(n); }); + if (f.tellg() != bytes) { + RAFT_LOG_WARN("Ignoring BBQ cache of unexpected size: %s", path.c_str()); + return std::nullopt; + } + f.seekg(0); + for_each_buffer( + q, [&f](void* p, size_t n) { f.read(static_cast(p), static_cast(n)); }); + if (!f) { + RAFT_LOG_WARN("Failed to read BBQ cache, re-quantizing: %s", path.c_str()); + return std::nullopt; + } + // Cheaper to recompute than to store and validate. dequant_delta/dequant_sum_delta are likewise + // a deterministic function of the cached lower/upper_intervals, quantized_component_sums, and + // layout, so they are recomputed here rather than added to the on-disk layout. + for (int64_t d = 0; d < dim; ++d) { + q.centroid_norm_sq += q.centroid(d) * q.centroid(d); + } + derive_dequant_factors(q); + RAFT_LOG_INFO("Loaded BBQ codes from %s", path.c_str()); + return q; +} + +/** Writes via a temporary so an interrupted run cannot leave a truncated cache behind. */ +inline void cache_store(const std::string& path, host_quantizer_storage& q) +{ + const std::string tmp = path + ".tmp"; + { + std::ofstream f(tmp, std::ios::binary | std::ios::trunc); + if (!f) { + RAFT_LOG_WARN("Cannot write BBQ cache: %s", tmp.c_str()); + return; + } + for_each_buffer(q, [&f](void* p, size_t n) { + f.write(static_cast(p), static_cast(n)); + }); + if (!f) { + RAFT_LOG_WARN("Failed to write BBQ cache: %s", tmp.c_str()); + f.close(); + std::remove(tmp.c_str()); + return; + } + } + if (std::rename(tmp.c_str(), path.c_str()) != 0) { + RAFT_LOG_WARN("Cannot move BBQ cache into place: %s", path.c_str()); + std::remove(tmp.c_str()); + return; + } + RAFT_LOG_INFO("Wrote BBQ codes to %s", path.c_str()); +} + +inline auto quantize_cached(const float* rows, + int64_t n_rows, + int64_t dim, + bbq_code_layout layout, + cuvs::distance::DistanceType metric) -> host_quantizer_storage +{ + const auto path = cache_path(n_rows, dim, layout, metric); + if (auto cached = cache_load(path, n_rows, dim, layout, metric)) { return std::move(*cached); } + auto quantized = quantize(rows, n_rows, dim, metric, layout); + cache_store(path, quantized); + return quantized; +} + +/** + * Quantize host-resident @p rows and upload the codes, ready for the BBQ NN-Descent build. + * + * `query_token == doc_token` produces a single quantizer and a symmetric join; differing tokens + * produce two quantizers and an asymmetric join, where the wider codes serve the query side. + */ +inline auto quantize_to_device(raft::resources const& res, + const float* rows, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + bbq_code_layout query_layout, + bbq_code_layout doc_layout) + -> cuvs::neighbors::device_bbq_dataset +{ + validate_layout_pair(query_layout, doc_layout); + std::vector host; + host.push_back(quantize_cached(rows, n_rows, dim, query_layout, metric)); + if (doc_layout != query_layout) { + host.push_back(quantize_cached(rows, n_rows, dim, doc_layout, metric)); + } + auto device = make_device_bbq_dataset(res, host); + // The uploads are stream-ordered against `host`, which dies with this frame. + raft::resource::sync_stream(res); + return device; +} + +inline auto quantize_to_device(raft::resources const& res, + const float* rows, + int64_t n_rows, + int64_t dim, + cuvs::distance::DistanceType metric, + std::string_view query_token, + std::string_view doc_token) + -> cuvs::neighbors::device_bbq_dataset +{ + return quantize_to_device(res, + rows, + n_rows, + dim, + metric, + parse_bbq_layout_token(query_token), + parse_bbq_layout_token(doc_token)); +} +} // namespace cuvs_internal::bbq diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index e2cefcd78d..0e1454d987 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -45,7 +45,8 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso if constexpr (nb::is_padded_dataset_view_v || nb::is_standard_dataset_view_v) { rm_dataset = dataset_.view(); - } else if constexpr (nb::is_vpq_dataset_view_v) { + } else if constexpr (nb::is_vpq_dataset_view_v || + nb::is_bbq_dataset_view_v) { skip_norms = true; } @@ -304,7 +305,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until the type-changing update_dataset(...) supplies a device-padded dataset. - if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + return cuvs::neighbors::cagra::detail::build_from_bbq_dataset( + res, params, dataset); + } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { auto effective_params = params; if (std::holds_alternative(effective_params.graph_build_params)) { effective_params.graph_build_params = graph_build_params::iterative_search_params{}; diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index b4c83c6191..2a3e2b43c1 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -20,6 +20,7 @@ using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_vie using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; using inst_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; } // namespace namespace cuvs::neighbors::cagra { @@ -32,6 +33,8 @@ extern template void index::compute_d raft::resources const&); extern template void index::compute_dataset_norms_( raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); void build_knn_graph(raft::resources const& handle, raft::host_matrix_view dataset, raft::host_matrix_view knn_graph, @@ -62,6 +65,9 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_vpq_view_t, cuvs::neighbors::cagra::device_pq_index); #endif +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_bbq_view_t, + cuvs::neighbors::cagra::device_bbq_index); + #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD #define CUVS_INST_CAGRA_UPDATE_DATASET(T, IdxT, SrcDatasetViewT, DstDatasetViewT) \ @@ -100,6 +106,8 @@ CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_device_standard_view_t, ins CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_device_padded_view_t, inst_vpq_view_t); CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_vpq_view_t, inst_device_padded_view_t); CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_vpq_view_t, inst_vpq_view_t); +CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_bbq_view_t, inst_device_padded_view_t); +CUVS_INST_CAGRA_UPDATE_DATASET(data_t, index_t, inst_bbq_view_t, inst_vpq_view_t); #undef CUVS_INST_CAGRA_UPDATE_DATASET diff --git a/cpp/src/neighbors/cagra_search_inst.cu.in b/cpp/src/neighbors/cagra_search_inst.cu.in index c5e35819f7..9fff49abd6 100644 --- a/cpp/src/neighbors/cagra_search_inst.cu.in +++ b/cpp/src/neighbors/cagra_search_inst.cu.in @@ -14,6 +14,7 @@ using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view< using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; using inst_vpq_f32_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; using inst_empty_view_t = cuvs::neighbors::device_empty_dataset_view; } // namespace @@ -57,6 +58,8 @@ index::compute_dataset_norms_(raft::resou template CUVS_EXPORT __attribute__((externally_visible)) void index::compute_dataset_norms_(raft::resources const& res); template CUVS_EXPORT __attribute__((externally_visible)) void +index::compute_dataset_norms_(raft::resources const& res); +template CUVS_EXPORT __attribute__((externally_visible)) void index::compute_dataset_norms_(raft::resources const& res); #define CUVS_INST_CAGRA_SEARCH_MULTI_PARTITION(T, IdxT, OutputIdxT) \ diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index c05e4c0deb..18e807b90f 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -184,6 +184,35 @@ namespace cuvs::neighbors::cagra { cuvs::neighbors::cagra::detail::deserialize(handle, is, index, nullptr); \ } \ \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::device_bbq_index& index) \ + { \ + cuvs::neighbors::cagra::detail::serialize(handle, filename, index, false); \ + } \ + \ + void deserialize(raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::device_bbq_index* index) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, nullptr); \ + } \ + \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::device_bbq_index& index) \ + { \ + cuvs::neighbors::cagra::detail::serialize(handle, os, index, false); \ + } \ + \ + void deserialize(raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::device_bbq_index* index) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, nullptr); \ + } \ + \ void serialize_to_hnswlib( \ raft::resources const& handle, \ std::ostream& os, \ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index d015236eca..9a22676334 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2212,6 +2212,36 @@ void build_knn_graph( res, build_params.metric, dataset, knn_graph_internal); } +template +void build_knn_graph(raft::resources const& res, + cuvs::neighbors::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph, + cuvs::neighbors::nn_descent::index_params build_params) +{ + raft::common::nvtx::range fun_scope( + "cagra::build_knn_graph(%zu, %zu, %u)", + size_t(dataset.n_rows()), + size_t(dataset.dim()), + size_t(knn_graph.extent(1))); + + std::optional> graph_view = knn_graph; + auto nn_descent_idx = cuvs::neighbors::nn_descent::build(res, build_params, dataset, graph_view); + + using internal_IdxT = typename std::make_unsigned::type; + using g_accessor = typename decltype(nn_descent_idx.graph())::accessor_type; + using g_accessor_internal = + raft::host_device_accessor, g_accessor::mem_type>; + + auto knn_graph_internal = + raft::mdspan, raft::row_major, g_accessor_internal>( + reinterpret_cast(nn_descent_idx.graph().data_handle()), + nn_descent_idx.graph().extent(0), + nn_descent_idx.graph().extent(1)); + + cuvs::neighbors::cagra::detail::graph::sort_knn_graph_bbq( + res, build_params.metric, dataset, knn_graph_internal); +} + template , raft::memory_type::host>, @@ -2871,4 +2901,94 @@ auto build_from_device_matrix(raft::resources const& res, RAFT_LOG_TRACE("Graph optimized, creating index"); return idx; } + +[[nodiscard]] inline auto resolve_bbq_knn_graph_build_params(index_params const& params, + size_t intermediate_degree) + -> cuvs::neighbors::nn_descent::index_params +{ + if (std::holds_alternative(params.graph_build_params)) { + return cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); + } + + auto nn_descent_params = + std::get(params.graph_build_params); + if (nn_descent_params.metric != params.metric) { + RAFT_LOG_WARN( + "Metric (%lu) for nn-descent needs to match cagra metric (%lu), " + "aligning nn-descent metric.", + nn_descent_params.metric, + params.metric); + nn_descent_params.metric = params.metric; + } + if (nn_descent_params.graph_degree != intermediate_degree) { + RAFT_LOG_WARN( + "Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), " + "aligning nn-descent graph_degree.", + nn_descent_params.graph_degree, + intermediate_degree); + nn_descent_params = + cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric); + } + return nn_descent_params; +} + +/** + * Build from a device-resident BBQ-quantized dataset: the whole graph construction runs on the + * compressed codes. + * + * The returned index cannot be searched, because CAGRA has no BBQ search kernels. Pass an + * uncompressed device-padded dataset to `cagra::update_dataset` to obtain a searchable index over + * the same graph. + */ +template + requires cuvs::neighbors::is_device_bbq_dataset_view_v +auto build_from_bbq_dataset(raft::resources const& res, + const index_params& params, + DatasetViewT const& dataset) + -> cuvs::neighbors::cagra::index +{ + RAFT_EXPECTS(!dataset.quantizers.empty(), "cagra::build: the BBQ dataset is empty."); + RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded || + params.metric == cuvs::distance::DistanceType::InnerProduct, + "cagra::build: a BBQ-quantized dataset supports L2Expanded, L2SqrtExpanded, " + "CosineExpanded, and InnerProduct."); + RAFT_EXPECTS(std::holds_alternative(params.graph_build_params) || + std::holds_alternative( + params.graph_build_params), + "cagra::build: a BBQ-quantized dataset requires nn-descent graph construction."); + + size_t intermediate_degree = params.intermediate_graph_degree; + size_t graph_degree = params.graph_degree; + common::nvtx::range function_scope( + "cagra::detail::build_from_bbq_dataset(%zu, %zu)", intermediate_degree, graph_degree); + auto const n_rows = static_cast(dataset.n_rows()); + check_graph_degree(intermediate_degree, graph_degree, static_cast(n_rows)); + + auto nn_descent_params = resolve_bbq_knn_graph_build_params(params, intermediate_degree); + nn_descent_params.return_distances = false; + + auto cagra_graph = [&]() -> raft::host_matrix { + std::optional> knn_graph( + raft::make_host_matrix(n_rows, intermediate_degree)); + build_knn_graph(res, dataset, knn_graph->view(), nn_descent_params); + + auto optimized = raft::make_host_matrix(n_rows, graph_degree); + RAFT_LOG_TRACE("optimizing graph"); + optimize(res, knn_graph->view(), optimized.view(), params.guarantee_connectivity); + knn_graph.reset(); + return optimized; + }(); + + RAFT_LOG_TRACE("Graph optimized, creating index"); + + if (params.attach_dataset_on_build) { + return cuvs::neighbors::cagra::index( + res, params.metric, dataset, raft::make_const_mdspan(cagra_graph.view())); + } + cuvs::neighbors::cagra::index idx(res, params.metric); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + return idx; +} } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 438303a607..94ea20c1ad 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -77,6 +77,19 @@ constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) return raw <= static_cast(kind::host_standard); } +/** + * Quantized datasets (PQ, BBQ) are owned outside the index and carry codebooks the index file has + * no representation for, so such indexes serialize the graph alone. + */ +template +inline constexpr bool is_graph_only_dataset_view_v = + cuvs::neighbors::is_vpq_dataset_view_v || + cuvs::neighbors::is_bbq_dataset_view_v; + +inline constexpr char const* kGraphOnlyDatasetMessage = + "CAGRA indexes with a quantized dataset store only the graph; serialize the quantized dataset " + "separately and reattach it with update_dataset()"; + template void serialize_index_mdspan(raft::resources const& res, std::ostream& os, MdspanT const& mdspan) { @@ -111,10 +124,8 @@ void serialize(raft::resources const& res, include_dataset &= (index_.dataset().n_rows() > 0); auto dataset_kind = cuvs::neighbors::cagra::serialized_dataset_kind::none; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_EXPECTS(!include_dataset, - "CAGRA PQ index serialization stores the graph only; serialize the PQ dataset " - "separately"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_EXPECTS(!include_dataset, kGraphOnlyDatasetMessage); } else { if (include_dataset) { dataset_kind = serialized_dataset_kind_for_view(); } } @@ -140,8 +151,8 @@ void serialize(raft::resources const& res, RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); - } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("CAGRA PQ index serialization stores the graph only"); + } else if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL(kGraphOnlyDatasetMessage); } else { // A further dataset type requires a new branch here and a corresponding deserialize branch. // Use static_assert to catch unsupported types at compile time. @@ -588,14 +599,14 @@ void deserialize_impl( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cagra::deserialize: PQ index files must contain only the graph"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL("cagra::deserialize: quantized index files must contain only the graph"); } else { cuvs::neighbors::detail::skip_dense_dataset(res, is); } } else { - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { - RAFT_FAIL("cagra::deserialize: PQ index files must contain only the graph"); + if constexpr (is_graph_only_dataset_view_v) { + RAFT_FAIL("cagra::deserialize: quantized index files must contain only the graph"); } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index cdcc18867b..ce522533a8 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -82,6 +82,7 @@ namespace descriptor_cache { */ struct key { uint64_t data_ptr; + uint64_t norms_ptr; uint64_t n_rows; uint32_t dim; uint32_t extra_val; // this one has different meanings for different descriptor types @@ -93,10 +94,11 @@ struct key { template auto make_key(const cagra::search_params& params, const DatasetT& dataset, - cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + cuvs::distance::DistanceType metric, + const void* dataset_norms) -> std::enable_if_t, key> { return key{reinterpret_cast(dataset.view().data_handle()), + reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), dataset.stride(), @@ -108,10 +110,11 @@ auto make_key(const cagra::search_params& params, template auto make_key(const cagra::search_params& params, const DatasetT& dataset, - cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + cuvs::distance::DistanceType metric, + const void* dataset_norms) -> std::enable_if_t, key> { return key{reinterpret_cast(dataset.data.data_handle()), + reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), uint32_t(reinterpret_cast(dataset.pq_code_book.data_handle()) >> 6), @@ -122,15 +125,16 @@ auto make_key(const cagra::search_params& params, inline auto operator==(const key& a, const key& b) -> bool { - return a.data_ptr == b.data_ptr && a.n_rows == b.n_rows && a.dim == b.dim && - a.extra_val == b.extra_val && a.team_size == b.team_size && a.metric == b.metric && - a.smem_dtype == b.smem_dtype; + return a.data_ptr == b.data_ptr && a.norms_ptr == b.norms_ptr && a.n_rows == b.n_rows && + a.dim == b.dim && a.extra_val == b.extra_val && a.team_size == b.team_size && + a.metric == b.metric && a.smem_dtype == b.smem_dtype; } struct key_hash { inline auto operator()(const key& x) const noexcept -> std::size_t { - return size_t{x.data_ptr} + size_t{x.n_rows} * size_t{x.dim} * size_t{x.extra_val} + + return size_t{x.data_ptr} + size_t{x.norms_ptr} + + size_t{x.n_rows} * size_t{x.dim} * size_t{x.extra_val} + (size_t{x.team_size} ^ size_t{x.metric}) + size_t{x.smem_dtype}; } }; @@ -164,7 +168,7 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, const DistanceT* dataset_norms = nullptr) -> dataset_descriptor_host { - auto key = descriptor_cache::make_key(params, dataset, metric); + auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms); auto& cache = raft::resource::get_custom_resource>(res) ->value; diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index c088b8c477..8dcb4d8b6f 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -8,13 +8,18 @@ #include "utils.hpp" // TODO: This shouldn't be invoking anything from spatial/knn +#include "../../../preprocessing/quantize/detail/bbq_distance.cuh" #include "../ann_utils.cuh" +#include +#include #include +#include #include #include #include +#include namespace cuvs::neighbors::cagra::detail::graph { namespace { @@ -156,6 +161,190 @@ void launch_sort_knn_graph_impl(raft::resources const& res, RAFT_CUDA_TRY(cudaGetLastError()); } +template +using bbq_quantizer_view_t = cuvs::preprocessing::quantize::bbq::quantizer_view; + +/** + * Distance between two dataset rows in quantized space + * + */ +template +__device__ __forceinline__ float bbq_row_distance( + const bbq_quantizer_view_t& quantizer_document, + const bbq_quantizer_view_t& quantizer_query, + cuvs::distance::DistanceType metric, + int64_t row_document, + int64_t row_query) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + const uint32_t raw = + bbq::code_inner_product(quantizer_document, quantizer_query, row_document, row_query); + return bbq::bbq_calculate_metric(raw, + bbq::get_dequant_factors(quantizer_document, row_document), + bbq::get_dequant_factors(quantizer_query, row_query), + quantizer_document, + quantizer_query, + metric, + raft::identity_op{}, + row_document, + row_query); +} + +template +__global__ void kern_sort_bbq(const bbq_quantizer_view_t quantizer_document, + const bbq_quantizer_view_t quantizer_query, + uint32_t* const knn_graph, // [graph_size, graph_degree] + const uint32_t graph_size, + const uint32_t graph_degree, + const cuvs::distance::DistanceType metric) +{ + const uint32_t src_node = (blockDim.x * blockIdx.x + threadIdx.x) / raft::WarpSize; + if (src_node >= graph_size) { return; } + + const uint32_t lane_id = threadIdx.x % raft::WarpSize; + + float my_keys[numElementsPerThread]; + uint32_t my_vals[numElementsPerThread]; + + for (int i = 0; i < numElementsPerThread; i++) { + const uint32_t k = i * raft::WarpSize + lane_id; + if (k >= graph_degree) { + my_keys[i] = utils::get_max_value(); + my_vals[i] = utils::get_max_value(); + continue; + } + const uint32_t dst_node = knn_graph[k + static_cast(graph_degree) * src_node]; + my_keys[i] = bbq_row_distance(quantizer_document, quantizer_query, metric, src_node, dst_node); + my_vals[i] = dst_node; + } + + raft::util::bitonic(true).sort(my_keys, my_vals); + + for (int i = 0; i < numElementsPerThread; i++) { + const uint32_t k = i * raft::WarpSize + lane_id; + if (k < graph_degree) { + knn_graph[k + (static_cast(graph_degree) * src_node)] = my_vals[i]; + } + } +} + +template +using sort_bbq_kernel_type = void (*)(bbq_quantizer_view_t, + bbq_quantizer_view_t, + uint32_t*, + uint32_t, + uint32_t, + cuvs::distance::DistanceType); + +template +auto select_sort_bbq_kernel(uint32_t degree) -> sort_bbq_kernel_type +{ + if (degree <= raft::WarpSize * 1) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 2) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 4) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 8) { return kern_sort_bbq; } + if (degree <= raft::WarpSize * 16) { return kern_sort_bbq; } + if (degree <= kMaxSortDegree) { return kern_sort_bbq; } + RAFT_FAIL( + "The degree of input knn graph is too large (%u). It must be equal to or smaller than %lu.", + degree, + kMaxSortDegree); +} + +template +auto select_sort_quantizers(cuvs::neighbors::device_bbq_dataset_view const& dataset) + -> std::pair, bbq_quantizer_view_t> +{ + using bbq_code_layout = cuvs::preprocessing::quantize::bbq::bbq_code_layout; + + if (dataset.quantizers.size() == 1) { return {dataset.quantizers[0], dataset.quantizers[0]}; } + + const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + + const bool tc_pair = has_1b && has_4b; + const bool simt_pair = (has_4bt && (has_1b || has_2bt)) || (has_2bt && has_1b); + RAFT_EXPECTS(tc_pair || simt_pair, + "Unsupported BBQ layout pair for an asymmetric dataset. Supported (document, " + "query) pairs: (packed_1b, packed_4b), (packed_1b, transposed_2b), " + "(packed_1b, transposed_4b), (transposed_2b, transposed_4b)."); + return {has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b), + tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) + : has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b)}; +} + +template +void sort_knn_graph_bbq_impl(raft::resources const& res, + cuvs::distance::DistanceType metric, + cuvs::neighbors::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + + RAFT_EXPECTS(!dataset.quantizers.empty(), "the BBQ dataset holds no quantizer"); + RAFT_EXPECTS(dataset.n_rows() == knn_graph.extent(0), + "dataset size is expected to have the same number of graph index size"); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded, + "Unsupported metric. Only L2Expanded, L2SqrtExpanded, InnerProduct and " + "CosineExpanded are supported for a BBQ-quantized dataset"); + + auto const graph_size = static_cast(knn_graph.extent(0)); + auto const graph_degree = static_cast(knn_graph.extent(1)); + auto kernel = select_sort_bbq_kernel(graph_degree); + auto const [quantizer_document, quantizer_query] = select_sort_quantizers(dataset); + // The code inner products read both rows as uint32_t words, and a bit-sliced layout starts + // every plane at a multiple of the plane stride, so each plane must be 4-byte aligned. + for (const auto& quantizer : {quantizer_document, quantizer_query}) { + const auto row_length = bbq::get_encoded_row_length(quantizer); + const auto planes = static_cast(bbq::get_code_planes(quantizer.layout)); + RAFT_EXPECTS(row_length % (4u * planes) == 0, + "Sorting a BBQ-quantized kNN graph requires the encoded row length to be a " + "multiple of 4*n_planes for 32-bit aligned plane loads, got %u with n_planes = %u", + row_length, + planes); + } + // A packed_1b document is promoted to 4-bit width one 32-dimension word at a time, which only + // covers the packed_4b query row exactly when the dimensionality is a multiple of 32. + RAFT_EXPECTS(quantizer_query.layout != bbq::bbq_code_layout::packed_4b || + quantizer_document.layout == quantizer_query.layout || + quantizer_document.dim() % 32 == 0, + "Sorting a BBQ-quantized kNN graph with packed_1b codes against packed_4b ones " + "requires the dataset dim to be a multiple of 32, got %u", + quantizer_document.dim()); + + const double time_sort_start = cur_time(); + RAFT_LOG_DEBUG("# Sorting kNN Graph on GPUs\n"); + + auto large_tmp_mr = raft::resource::get_large_workspace_resource_ref(res); + auto d_input_graph = raft::make_device_mdarray( + res, large_tmp_mr, raft::make_extents(graph_size, graph_degree)); + raft::copy(res, d_input_graph.view(), knn_graph); + + constexpr uint32_t block_size = 256; + auto const warps = block_size / raft::WarpSize; + auto const blocks = (graph_size + warps - 1) / warps; + kernel<<>>( + quantizer_document, + quantizer_query, + d_input_graph.data_handle(), + graph_size, + graph_degree, + metric); + RAFT_CUDA_TRY(cudaGetLastError()); + raft::resource::sync_stream(res); + raft::copy(res, knn_graph, raft::make_const_mdspan(d_input_graph.view())); + + const double time_sort_end = cur_time(); + RAFT_LOG_DEBUG("# Sorting kNN graph time: %.1lf sec\n", time_sort_end - time_sort_start); +} + } // namespace #define CUVS_DEFINE_CAGRA_GRAPH_SORT(DataT) \ @@ -178,6 +367,22 @@ CUVS_DEFINE_CAGRA_GRAPH_SORT(uint8_t) #undef CUVS_DEFINE_CAGRA_GRAPH_SORT +#define CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(DataT) \ + void sort_knn_graph_bbq(raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::device_bbq_dataset_view dataset, \ + raft::host_matrix_view knn_graph) \ + { \ + sort_knn_graph_bbq_impl(res, metric, dataset, knn_graph); \ + } + +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(float) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(half) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(int8_t) +CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(uint8_t) + +#undef CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ + void optimize_device_graph( raft::resources const& res, raft::device_matrix_view knn_graph, diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cuh b/cpp/src/neighbors/detail/cagra/graph_shared.cuh index 0f2bc0f7dc..33ec7f92ed 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cuh @@ -6,8 +6,10 @@ #include #include +#include #include +#include #include #include @@ -35,6 +37,20 @@ CUVS_DECL_CAGRA_GRAPH_SORT(uint8_t); #undef CUVS_DECL_CAGRA_GRAPH_SORT +#define CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(DataT) \ + CUVS_EXPORT void sort_knn_graph_bbq( \ + raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::device_bbq_dataset_view dataset, \ + raft::host_matrix_view knn_graph) + +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(float); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(half); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(int8_t); +CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(uint8_t); + +#undef CUVS_DECL_CAGRA_GRAPH_SORT_BBQ + /** Run the existing CAGRA optimizer through one compiled instantiation instead of rematerializing * its reverse-graph, prune, merge, and MST kernels in every Fastener dtype TU. */ CUVS_EXPORT void optimize_device_graph( diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index ec92e1a36f..b1260784b1 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -9,9 +9,12 @@ #include "neighbors_device_intrinsics.cuh" #include "nn_descent_gnnd.hpp" +#include "../../core/nvtx.hpp" #include "../../core/omp_wrapper.hpp" +#include "../../preprocessing/quantize/detail/bbq_distance.cuh" #include #include +#include #include #include @@ -38,6 +41,7 @@ #include +#include #include #include #include @@ -47,6 +51,9 @@ namespace cuvs::neighbors::nn_descent::detail { +using cuvs::preprocessing::quantize::bbq::bbq_code_layout; +using cuvs::preprocessing::quantize::bbq::quantizer_view; + template struct ResultItem; @@ -423,7 +430,9 @@ __device__ ResultItem get_min_item(const Index_t id, const int idx_in_list, const Index_t* neighbs, const DistData_t* distances, - const bool find_in_row = true) + const bool find_in_row = true, + const int stride = SKEWED_MAX_NUM_BI_SAMPLES, + const int neighbs_size = MAX_NUM_BI_SAMPLES) { int lane_id = threadIdx.x % raft::warp_size(); @@ -434,16 +443,18 @@ __device__ ResultItem get_min_item(const Index_t id, idx[0] = lane_id; idx[1] = raft::warp_size() + lane_id; - if (neighbs[idx[0]] != id) { - dist[0] = find_in_row ? distances[idx_in_list * SKEWED_MAX_NUM_BI_SAMPLES + lane_id] - : distances[idx_in_list + lane_id * SKEWED_MAX_NUM_BI_SAMPLES]; + // neighbs_size defaults to the full width, so existing (dense/calculate_metric-based) callers + // that pre-fill `distances` out to MAX_NUM_BI_SAMPLES are unaffected. Callers that don't pre-fill + // (e.g. the BBQ SIMT kernel, which only ever writes real cells) pass the real list size instead, + // so out-of-range entries just never get read -- cheaper than a separate fill pass. + if (idx[0] < neighbs_size && neighbs[idx[0]] != id) { + dist[0] = find_in_row ? distances[idx_in_list * stride + lane_id] + : distances[idx_in_list + lane_id * stride]; } - if (neighbs[idx[1]] != id) { - dist[1] = - find_in_row - ? distances[idx_in_list * SKEWED_MAX_NUM_BI_SAMPLES + raft::warp_size() + lane_id] - : distances[idx_in_list + (raft::warp_size() + lane_id) * SKEWED_MAX_NUM_BI_SAMPLES]; + if (idx[1] < neighbs_size && neighbs[idx[1]] != id) { + dist[1] = find_in_row ? distances[idx_in_list * stride + raft::warp_size() + lane_id] + : distances[idx_in_list + (raft::warp_size() + lane_id) * stride]; } if (dist[1] < dist[0]) { @@ -523,7 +534,7 @@ __device__ __forceinline__ void calculate_metric(float* s_distances, { // if we have a distance epilogue, distances need to be fully calculated instead of postprocessing // them. - bool can_postprocess_dist = std::is_same_v; + constexpr bool can_postprocess_dist = std::is_same_v; for (int i = threadIdx.x; i < MAX_NUM_BI_SAMPLES * SKEWED_MAX_NUM_BI_SAMPLES; i += blockDim.x) { int row_id = i / SKEWED_MAX_NUM_BI_SAMPLES; @@ -583,6 +594,11 @@ struct DistAccumulator { // SIMT kernel: scalar element-wise distance computation. // Used for fp32 data (all metrics) and L1 distance computation for all dtypes (which cannot use // tensor cores). +// -------------------------------------------------------------------------- +// Dense local-join kernels +// Full-precision joins over the original vectors: SIMT and fp16 tensor core. +// -------------------------------------------------------------------------- + template , @@ -848,13 +864,6 @@ __launch_bounds__(BLOCK_SIZE) #endif } -// launch_bounds here denote BLOCK_SIZE = 512 and MIN_BLOCKS_PER_SM = 4 -// Per -// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications, -// MAX_RESIDENT_THREAD_PER_SM = BLOCK_SIZE * BLOCKS_PER_SM = 2048 -// For architectures 750 and 860 (890), the values for MAX_RESIDENT_THREAD_PER_SM -// is 1024 and 1536 respectively, which means the bounds don't work anymore -// Used for fp32 data downcast to fp16, and all types using non-L1 distance metric. template , @@ -1115,6 +1124,905 @@ __launch_bounds__(BLOCK_SIZE) } #endif } +// -------------------------------------------------------------------------- +// BBQ local-join kernels +// Quantized-code joins: SIMT (popc / dp4a) and int4 tensor core (u4 wmma). +// -------------------------------------------------------------------------- + +// The code promotion, the dequant factors and the raw-dot-product-to-distance conversion live +// next to the code inner products in bbq.cuh, so the CAGRA graph sort ranks by the very same +// distance these kernels build the lists with. +using cuvs::preprocessing::quantize::bbq::bbq_calculate_metric; +using cuvs::preprocessing::quantize::bbq::bbq_dequant_factors; +using cuvs::preprocessing::quantize::bbq::get_dequant_factors; +using cuvs::preprocessing::quantize::bbq::packed_1b_to_4b; + +// Stages one K-tile of `count` neighbor rows into SMEM. Shared skeleton: one warp per row, +// lanes strided along native uint32 words, last-tile zero-pad so compute reads a full tile. +// +// stage_tile_simt copies native codes in storage layout (Planes > 1 gathers bit-sliced planes). +// stage_promoted_tile expands packed_1b to u4 for int4 MMA (packed_4b is a plain copy). +template +__device__ __forceinline__ void stage_tile_simt(uint8_t (*dst)[RowStride], + const quantizer_view& quantizer, + const Index_t* neighbors, + const int count, + const size_t base, + const int plane_extent, + const int num_load_u32, + const int plane_tile_u32, + const bool last_tile, + const int warp_id, + const int lane_id) +{ + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + for (int i = 0; i < MAX_NUM_BI_SAMPLES / num_warps; ++i) { + const int idx = i * num_warps + warp_id; + if (idx >= count) continue; + auto* s = reinterpret_cast(dst[idx]); + const auto* src = reinterpret_cast(&quantizer.codes(neighbors[idx], base)); + for (int w = lane_id; w < num_load_u32; w += raft::warp_size()) { +#pragma unroll + for (int p = 0; p < Planes; ++p) { + s[p * plane_tile_u32 + w] = src[p * plane_extent + w]; + } + } + if (last_tile) { + for (int w = num_load_u32 + lane_id; w < plane_tile_u32; w += raft::warp_size()) { +#pragma unroll + for (int p = 0; p < Planes; ++p) { + s[p * plane_tile_u32 + w] = 0; + } + } + } + } +} + +// `native_row_bytes` is the layout's own encoded row length. One native tile is +// RowStride-independent: BBQ_ROW_BYTES / expansion native bytes promote to exactly BBQ_ROW_BYTES +// promoted bytes, so a single n_tiles drives every operand regardless of how compact each one's +// on-disk format is. +template +__device__ __forceinline__ void stage_promoted_tile(uint8_t (*dst)[RowStride], + const quantizer_view& quantizer, + const Index_t* neighbors, + const int count, + const int step, + const int native_row_bytes, + const int warp_id, + const int lane_id) +{ + static_assert(Layout == bbq_code_layout::packed_1b || Layout == bbq_code_layout::packed_4b, + "int4 MMA path supports packed_1b (1b), packed_4b (4b)"); + // A u4 MMA fragment needs 4 bits per value, so a layout storing `bits` bits per value expands + // one native word into 4/bits promoted words. packed_4b is the identity case (a plain word + // copy), which is why the symmetric kernel needs no separate "no promotion" path -- and why a + // promoted query works exactly like a promoted document. + constexpr int expansion = Layout == bbq_code_layout::packed_1b ? 4 : 1; + constexpr int native_tile = TileBytes / expansion; + constexpr int native_tile_u32 = native_tile / 4; + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + // packed_4b's native tile is 32 u32s (one per lane). packed_1b's is 8, so a warp-per-row + // map would leave 24 lanes idle; instead pack warp_size/8 = 4 rows into one pass. + static_assert(raft::warp_size() % native_tile_u32 == 0); + constexpr int rows_per_pass = raft::warp_size() / native_tile_u32; + static_assert(MAX_NUM_BI_SAMPLES % (num_warps * rows_per_pass) == 0); + + const int base = step * native_tile; + const int remaining = native_row_bytes - base; + const int num_load_u32 = + (remaining < native_tile ? (remaining > 0 ? remaining : 0) : native_tile) / 4; + + for (int i = 0; i < MAX_NUM_BI_SAMPLES / (num_warps * rows_per_pass); ++i) { + const int idx = (i * num_warps + warp_id) * rows_per_pass + lane_id / native_tile_u32; + const int w = lane_id % native_tile_u32; + if (idx >= count) continue; + auto* s = reinterpret_cast(dst[idx]); + const auto* src = reinterpret_cast(&quantizer.codes(neighbors[idx], base)); + if (w < num_load_u32) { + if constexpr (expansion == 4) { + reinterpret_cast(s)[w] = packed_1b_to_4b(src[w]); + } else { + s[w] = src[w]; + } + } else { + if constexpr (expansion == 4) { + reinterpret_cast(s)[w] = uint4{0, 0, 0, 0}; + } else { + s[w] = 0; + } + } + } +} + +// Stages this block's new/old neighbour lists into shared memory and drops duplicates. Identical +// in the SIMT and wmma BBQ kernels, so it lives here rather than twice. The caller keeps the +// `if (!new_size) return;` early-out, which has to stay at kernel scope. +template +__device__ __forceinline__ void stage_neighbor_lists(Index_t* new_neighbors, + Index_t* old_neighbors, + int* s_unique_counter, + const Index_t* graph_new, + const Index_t* rev_graph_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const size_t list_id, + const int width, + const int2 new_size2, + const int2 old_size2, + int& new_size, + int& old_size) +{ + const int tx = threadIdx.x; + if (tx < new_size2.x) { + new_neighbors[tx] = graph_new[list_id * width + tx]; + } else if (tx < new_size) { + new_neighbors[tx] = rev_graph_new[list_id * width + tx - new_size2.x]; + } + if (tx < old_size2.x) { + old_neighbors[tx] = graph_old[list_id * width + tx]; + } else if (tx < old_size) { + old_neighbors[tx] = rev_graph_old[list_id * width + tx - old_size2.x]; + } + __syncthreads(); + + remove_duplicates( + new_neighbors, new_size2.x, new_neighbors + new_size2.x, new_size2.y, s_unique_counter[0], 0); + remove_duplicates( + old_neighbors, old_size2.x, old_neighbors + old_size2.x, old_size2.y, s_unique_counter[1], 1); + __syncthreads(); + new_size = new_size2.x + s_unique_counter[0]; + old_size = old_size2.x + s_unique_counter[1]; +} + +template , + typename DistEpilogue_t> +RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) + local_join_kernel_bbq_simt(const Index_t* graph_new, + const Index_t* rev_graph_new, + const int2* sizes_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const int2* sizes_old, + const int width, + quantizer_view dataset_document, + quantizer_view dataset_query, + ID_t* graph, + DistData_t* dists, + int graph_width, + int* locks, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue) +{ + constexpr int document_planes = + cuvs::preprocessing::quantize::bbq::get_code_planes(DocumentLayout); + constexpr int query_planes = cuvs::preprocessing::quantize::bbq::get_code_planes(QueryLayout); + static_assert(!SelfJoin || DocumentLayout == QueryLayout, + "a self-join must use the same layout on both operands"); + + // Both operands are tiled at the same per-plane tile so each step covers the same dimension + // range on both sides. QUERY_ROW_BYTES fixes the query row width; the document row width then + // follows as QUERY_ROW_BYTES / (query_planes / document_planes), i.e. query_plane_tile scaled + // by the document's own plane count. Worked out per supported pair: + // + // pair query_plane_tile doc_row_bytes doc_stride query_stride + // ------ ---------------- -------------- ---------- ------------ + // 1 x 1 128 128 * 1 = 128 128 128 + // 2t x 2t 64 64 * 2 = 128 64 64 + // 1 + 2t 64 64 * 1 = 64 64 64 + // 1 + 4t 32 32 * 1 = 32 32 32 + // 2t + 4t 32 32 * 2 = 64 32 32 + // + constexpr int QUERY_ROW_BYTES = 128; + constexpr int BBQ_PAD = alignof(uint32_t); + // The document buffer is normally only the A operand (two rows broadcast across a warp), so it + // needs no bank-conflict pad. Under SelfJoin it doubles as the B operand in phase 1 (32 + // consecutive columns at one byte offset), so it needs the same skew the query buffer gets. + constexpr int DOC_PAD = SelfJoin ? BBQ_PAD : 0; + static_assert((QUERY_ROW_BYTES + BBQ_PAD) % alignof(uint32_t) == 0); + + __shared__ int s_list[MAX_NUM_BI_SAMPLES * 2]; + __shared__ __align__(alignof(uint32_t)) uint8_t + s_doc_vec[MAX_NUM_BI_SAMPLES][QUERY_ROW_BYTES / query_planes * document_planes + DOC_PAD]; + __shared__ __align__(alignof(uint32_t)) + uint8_t s_query_vec[MAX_NUM_BI_SAMPLES][QUERY_ROW_BYTES + BBQ_PAD]; + // Holds the final (post-metric, post-epilogue) float distance per cell + __shared__ DistData_t s_distances[MAX_NUM_BI_SAMPLES * SKEWED_MAX_NUM_BI_SAMPLES]; + __shared__ int s_unique_counter[2]; + // Document-side (row axis) dequant factors, indexed by list position -- shared by both phases + // below since rows are always the `new_neighbors`/document side in both, so this is staged once + // and never restaged. At a given accumulation step, up to MAX_NUM_BI_SAMPLES threads share the + // same row0 (see the pair_idx assignment below), so caching avoids up to 64x redundant scattered + // global reads + __shared__ bbq_dequant_factors s_document_factors[MAX_NUM_BI_SAMPLES]; + + if (threadIdx.x == 0) { + s_unique_counter[0] = 0; + s_unique_counter[1] = 0; + } + + Index_t* new_neighbors = s_list; + Index_t* old_neighbors = s_list + MAX_NUM_BI_SAMPLES; + const size_t list_id = blockIdx.x; + const int2 new_size2 = sizes_new[list_id]; + const int2 old_size2 = sizes_old[list_id]; + int new_size = new_size2.x + new_size2.y; + int old_size = old_size2.x + old_size2.y; + const int tx = threadIdx.x; + + if (!new_size) return; + stage_neighbor_lists(new_neighbors, + old_neighbors, + s_unique_counter, + graph_new, + rev_graph_new, + graph_old, + rev_graph_old, + list_id, + width, + new_size2, + old_size2, + new_size, + old_size); + + const int warp_id = threadIdx.x / raft::warp_size(); + const int lane_id = threadIdx.x % raft::warp_size(); + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + + // Rows are always new_neighbors/document in both phases below, so this runs exactly once. + for (int i = tx; i < new_size; i += BLOCK_SIZE) { + s_document_factors[i] = get_dequant_factors(dataset_document, new_neighbors[i]); + } + __syncthreads(); + + // Each plane gets an equal slice of the row in shared memory, so the cached bytes always form a + // valid encoded chunk. + // Bytes per plane = encoded row length / plane count. Do NOT assume ceildiv(dim, 8): that is + // bytes-per-plane only for the bit-plane layouts (packed_1b, transposed_2b, transposed_4b), + // where it happens to equal encoded/planes for all three. The dense byte layouts + // (packed_7b/packed_8b) are `dim` bytes in one plane, and would read 1/8 of each row. + const int plane_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_document)) / + document_planes; + assert(plane_bytes == static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length( + dataset_query)) / + query_planes); + constexpr int query_plane_tile = QUERY_ROW_BYTES / query_planes; + constexpr int plane_tile = query_plane_tile; + constexpr int doc_row_bytes = query_plane_tile * document_planes; + static_assert(plane_tile % 4 == 0, "plane_tile must be 4-byte aligned for uint32 loads"); + // Row strides too: rows are indexed as base + idx * stride and then read as uint32_t, so a + // stride that is not a multiple of 4 misaligns every odd row. Derived from QUERY_ROW_BYTES, so + // this is what catches an ill-chosen QUERY_ROW_BYTES rather than letting it fault at runtime. + static_assert((doc_row_bytes + DOC_PAD) % alignof(uint32_t) == 0, + "document row stride must be 4-byte aligned for uint32 loads"); + static_assert( + doc_row_bytes % document_planes == 0 && doc_row_bytes / document_planes == query_plane_tile, + "document plane stride must match query plane stride"); + // plane_bytes is the per-plane stride in bytes; plane_extent is the same in uint32 elements, + // computed once so call sites don't re-derive it. Alignment (plane_bytes % 4 == 0, i.e. + // dataset dim % 32 == 0) is enforced by the launcher. + const int plane_extent = plane_bytes / 4; + constexpr int plane_tile_u32 = plane_tile / 4; + constexpr int query_plane_tile_u32 = query_plane_tile / 4; + const int n_tiles = raft::ceildiv(plane_bytes, plane_tile); + + // with NUM_SAMPLES=32, BLOCK_SIZE=256, pairs_per_thread = 32 * 64 / 512 = 4 + constexpr int num_row_pairs = MAX_NUM_BI_SAMPLES / 2; + constexpr int num_pairs = num_row_pairs * MAX_NUM_BI_SAMPLES; + constexpr int pairs_per_thread = (num_pairs + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Every thread's column is invariant across its pairs_per_thread iterations, so fetch it once + // here and reuse the register copy in both phases' store loops below instead of a shared array. + const int my_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_col_factors{}; + if (my_col < new_size) { + my_col_factors = get_dequant_factors(dataset_query, new_neighbors[my_col]); + } + + uint32_t acc0[pairs_per_thread] = {}; + uint32_t acc1[pairs_per_thread] = {}; + for (int step = 0; step < n_tiles; ++step) { + const bool last_tile = (step == n_tiles - 1); + const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; + const int num_load_u32 = num_load / 4; + const size_t base = static_cast(step) * plane_tile; + stage_tile_simt(s_doc_vec, + dataset_document, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + plane_tile_u32, + last_tile, + warp_id, + lane_id); + + // Query and document tiles cover the same dimension range per step (both tile at + // query_plane_tile), so load the query tile once and run the dot product directly -- no + // per-step query sub-tile loop. Under SelfJoin phase 1 is new x new on a single quantizer, + // so the document buffer already holds exactly what the query buffer would: skip the load + // and point the B operand at s_doc_vec. + if constexpr (!SelfJoin) { + stage_tile_simt(s_query_vec, + dataset_query, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + query_plane_tile_u32, + last_tile, + warp_id, + lane_id); + } + __syncthreads(); + +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= new_size) continue; + // Phase 1 B operand: s_doc_vec under SelfJoin (see the staging note above). + const uint8_t* row_b; + if constexpr (SelfJoin) { + row_b = s_doc_vec[col]; + } else { + row_b = s_query_vec[col]; + } + cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( + s_doc_vec[row0], s_doc_vec[row0 + 1], row_b, acc0[k], acc1[k]); + } + __syncthreads(); + } +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= new_size) continue; + const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; + const Index_t query_id = new_neighbors[col]; + // row0 (unlike row0 + 1 below) isn't otherwise bounded by new_size -- gate it explicitly so a + // stale/garbage new_neighbors[row0] past the real list never reaches a dataset_document lookup + // (row_norm, for CosineExpanded); the cell is provably never read downstream either way, since + // the min-search loops below only visit idx_in_list < new_size. + if (row0 < new_size) { + const Index_t doc_id0 = new_neighbors[row0]; + s_distances[distance0] = bbq_calculate_metric(acc0[k], + s_document_factors[row0], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id0, + query_id); + } + if (row0 + 1 < new_size) { + const Index_t doc_id1 = new_neighbors[row0 + 1]; + s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = + bbq_calculate_metric(acc1[k], + s_document_factors[row0 + 1], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id1, + query_id); + } + } + __syncthreads(); + + for (int step = 0; step < raft::ceildiv(new_size, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + new_neighbors, + s_distances, + true, + SKEWED_MAX_NUM_BI_SAMPLES, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + if (!old_size) return; + __syncthreads(); + + // Phase 2's column list is old_neighbors, not new_neighbors -- refetch this thread's column + // factors (rows/s_document_factors stay valid unchanged, since rows are new_neighbors in both + // phases). + const int my_old_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_old_col_factors{}; + if (my_old_col < old_size) { + my_old_col_factors = get_dequant_factors(dataset_query, old_neighbors[my_old_col]); + } + + uint32_t acc0_old[pairs_per_thread] = {}; + uint32_t acc1_old[pairs_per_thread] = {}; + for (int step = 0; step < n_tiles; ++step) { + const bool last_tile = (step == n_tiles - 1); + const int num_load = last_tile ? plane_bytes - step * plane_tile : plane_tile; + const int num_load_u32 = num_load / 4; + const size_t base = static_cast(step) * plane_tile; + if (n_tiles > 1) { + stage_tile_simt(s_doc_vec, + dataset_document, + new_neighbors, + new_size, + base, + plane_extent, + num_load_u32, + plane_tile_u32, + last_tile, + warp_id, + lane_id); + } + stage_tile_simt(s_query_vec, + dataset_query, + old_neighbors, + old_size, + base, + plane_extent, + num_load_u32, + query_plane_tile_u32, + last_tile, + warp_id, + lane_id); + __syncthreads(); + +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= old_size) continue; + cuvs::preprocessing::quantize::bbq::bbq_code_inner_product_2x1( + s_doc_vec[row0], s_doc_vec[row0 + 1], s_query_vec[col], acc0_old[k], acc1_old[k]); + } + __syncthreads(); + } +#pragma unroll + for (int k = 0; k < pairs_per_thread; ++k) { + const int pair_idx = tx + k * BLOCK_SIZE; + if (pair_idx >= num_pairs) continue; + const int row0 = (pair_idx / MAX_NUM_BI_SAMPLES) * 2; + const int col = pair_idx % MAX_NUM_BI_SAMPLES; + if (col >= old_size) continue; + const int distance0 = row0 * SKEWED_MAX_NUM_BI_SAMPLES + col; + const Index_t query_id = old_neighbors[col]; + // See the identical row0 < new_size gate in phase 1 above -- same reasoning (row0 isn't + // otherwise bounded, and this cell is never read downstream either way). + if (row0 < new_size) { + const Index_t doc_id0 = new_neighbors[row0]; + s_distances[distance0] = bbq_calculate_metric(acc0_old[k], + s_document_factors[row0], + my_old_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id0, + query_id); + } + if (row0 + 1 < new_size) { + const Index_t doc_id1 = new_neighbors[row0 + 1]; + s_distances[distance0 + SKEWED_MAX_NUM_BI_SAMPLES] = + bbq_calculate_metric(acc1_old[k], + s_document_factors[row0 + 1], + my_old_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id1, + query_id); + } + } + __syncthreads(); + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + old_neighbors, + s_distances, + true, + SKEWED_MAX_NUM_BI_SAMPLES, + old_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= old_size) continue; + const int list_idx = idx_in_list + MAX_NUM_BI_SAMPLES; + auto min_elem = get_min_item(s_list[list_idx], + idx_in_list, + new_neighbors, + s_distances, + false, + SKEWED_MAX_NUM_BI_SAMPLES, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[list_idx], graph, dists, graph_width, locks); + } + } +} + +// int4 tensor-core BBQ local join, covering both the symmetric (one quantizer, self-join) and +// asymmetric (two quantizers) cases. Modeled on local_join_kernel_wmma: nvcuda::wmma fragments +// (u4 x u4 -> s32, shape m8n8k32) replace the popc/dp4a inner product; the accumulator lives in +// registers across the whole K reduction and is stored to s_distances once per (row-tile, +// col-tile), not accumulated into shared memory every step like the scalar kernel. +// +// Warp tiling: num_warps = BLOCK_SIZE/32 warps arranged as a WARPS_PER_DIM x WARPS_PER_DIM square +// grid (WARPS_PER_DIM=4 so 4x4=16=num_warps), each warp owning a (MAX_NUM_BI_SAMPLES/WARPS_PER_DIM) +// region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES output matrix, same as +// local_join_kernel_wmma's WMMA_M=N=16 warp assignment. Since int4 MMA tiles are MMA_M x MMA_N +// (8x8, the only shape nvcuda::wmma exposes for u4), each warp covers its region via a +// SUB_PER_DIM x SUB_PER_DIM grid of native tiles instead of a single call -- SUB_PER_DIM is a +// forced consequence of (MAX_NUM_BI_SAMPLES/MMA_M) / WARPS_PER_DIM, not an arbitrary choice. +template , + typename DistEpilogue_t> +RAFT_KERNEL __launch_bounds__(BLOCK_SIZE) + local_join_kernel_bbq_wmma(const Index_t* graph_new, + const Index_t* rev_graph_new, + const int2* sizes_new, + const Index_t* graph_old, + const Index_t* rev_graph_old, + const int2* sizes_old, + const int width, + const quantizer_view dataset_document, + const quantizer_view dataset_query, + ID_t* graph, + DistData_t* dists, + int graph_width, + int* locks, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue) +{ +// int4 sub-byte MMA (nvcuda::wmma experimental::precision::u4) still compiles on every Blackwell +// variant (sm_100/103/110/120/121, verified by disassembly): ptxas lowers it to the same software +// path it's always used since Turing -- unpack each u4 nibble pair into two u8 operands and run +// two native u8 IMMA instructions, summing the partial products. +#if (__CUDA_ARCH__ >= 750) + using namespace nvcuda; + constexpr int MMA_M = 8; + constexpr int MMA_N = 8; + constexpr int MMA_K = 32; + // num_warps = BLOCK_SIZE/32 = 16, arranged as a square WARPS_PER_DIM x WARPS_PER_DIM grid since + // 4*4=16 matches exactly; the static_assert is what actually enforces this holds for the + // current BLOCK_SIZE, WARPS_PER_DIM itself isn't derived (no trivial constexpr integer sqrt). + constexpr int WARPS_PER_DIM = 4; + static_assert(WARPS_PER_DIM * WARPS_PER_DIM == BLOCK_SIZE / raft::warp_size(), + "warp grid must be square and match num_warps = BLOCK_SIZE/32"); + // Each warp owns a WARP_TILE x WARP_TILE region of the MAX_NUM_BI_SAMPLES x MAX_NUM_BI_SAMPLES + // output matrix. TILES_PER_DIM is how many native MMA_M x MMA_N tiles span one output dimension; + // SUB_PER_DIM (native tiles per warp per dim) is a forced consequence of TILES_PER_DIM / + // WARPS_PER_DIM, not an arbitrary choice -- it's 2 here only because 8/4=2 for these particular + // MAX_NUM_BI_SAMPLES/MMA_M/WARPS_PER_DIM values. + static_assert(MAX_NUM_BI_SAMPLES % MMA_M == 0 && MMA_M == MMA_N, + "MAX_NUM_BI_SAMPLES must divide evenly into square MMA_MxMMA_N tiles"); + constexpr int TILES_PER_DIM = MAX_NUM_BI_SAMPLES / MMA_M; + static_assert(TILES_PER_DIM % WARPS_PER_DIM == 0, + "warps must evenly tile the native MMA tiles in each output dimension"); + constexpr int SUB_PER_DIM = TILES_PER_DIM / WARPS_PER_DIM; + constexpr int WARP_TILE = SUB_PER_DIM * MMA_M; + + // Promoted (4-bit-width) staging tile, 128 B/row. Row stride is BBQ_ROW_BYTES + MMA_PAD, not + // just BBQ_ROW_BYTES: sub-byte IMMA loads need at least 16-byte row alignment, and MMA_PAD must + // be a multiple of 16 to preserve that -- but BBQ_ROW_BYTES=128 alone is *also* exactly 32 + // shared-memory banks (4 B/bank), so every row would land on the same bank offset and any + // multi-row access load_matrix_sync does internally would conflict. MMA_PAD=16 breaks that + // exact-32-bank alignment (144 B/row is not a multiple of 128 B) while staying a multiple of 16 + // for the IMMA alignment requirement. + constexpr int BBQ_ROW_BYTES = 128; + constexpr int MMA_PAD = 16; + static_assert(MMA_PAD % 16 == 0, "row padding must preserve 16-byte IMMA row alignment"); + constexpr int ELEMS_PER_TILE = BBQ_ROW_BYTES * 2; // 2 u4 elements/byte + constexpr int K_STEPS_PER_TILE = ELEMS_PER_TILE / MMA_K; + constexpr int ROW_STRIDE_U4 = (BBQ_ROW_BYTES + MMA_PAD) * 2; // row-to-row stride, u4 elements + + constexpr int MMA_STORE_STRIDE = SKEWED_MAX_NUM_BI_SAMPLES; + + // s_row_vec is the A operand: always the `new` list, document quantizer. s_col_vec is the B + // operand: the `new` list in phase 1 and the `old` list in phase 2, query quantizer. Under + // SelfJoin, phase 1 leaves s_col_vec untouched and both fragments read s_row_vec. + __shared__ int s_list[MAX_NUM_BI_SAMPLES * 2]; + __shared__ __align__(16) uint8_t s_row_vec[MAX_NUM_BI_SAMPLES][BBQ_ROW_BYTES + MMA_PAD]; + // s_col_vec aliases s_distances's memory instead of its own array: their live ranges never + // overlap (col_vec fully consumed by the kk-loop before distances are stored; distances fully + // consumed by the min-search loop before the next phase re-stages col_vec). + __shared__ __align__(16) DistData_t s_distances[MAX_NUM_BI_SAMPLES * MMA_STORE_STRIDE]; + static_assert(sizeof(s_distances) >= MAX_NUM_BI_SAMPLES * (BBQ_ROW_BYTES + MMA_PAD), + "s_col_vec aliases s_distances's memory and must fit inside it"); + auto(*s_col_vec)[BBQ_ROW_BYTES + MMA_PAD] = + reinterpret_cast(s_distances); + __shared__ int s_unique_counter[2]; + // Document-side (row axis) dequant factors -- see the identical buffer and full rationale in + // local_join_kernel_bbq_simt. Staged once below, reused unchanged by both phases (rows are + // always new_neighbors/document in both). + __shared__ bbq_dequant_factors s_document_factors[MAX_NUM_BI_SAMPLES]; + + if (threadIdx.x == 0) { + s_unique_counter[0] = 0; + s_unique_counter[1] = 0; + } + + Index_t* new_neighbors = s_list; + Index_t* old_neighbors = s_list + MAX_NUM_BI_SAMPLES; + const size_t list_id = blockIdx.x; + const int2 new_size2 = sizes_new[list_id]; + const int2 old_size2 = sizes_old[list_id]; + int new_size = new_size2.x + new_size2.y; + int old_size = old_size2.x + old_size2.y; + const int tx = threadIdx.x; + + if (!new_size) return; + stage_neighbor_lists(new_neighbors, + old_neighbors, + s_unique_counter, + graph_new, + rev_graph_new, + graph_old, + rev_graph_old, + list_id, + width, + new_size2, + old_size2, + new_size, + old_size); + + const int warp_id = threadIdx.x / raft::warp_size(); + const int lane_id = threadIdx.x % raft::warp_size(); + constexpr int num_warps = BLOCK_SIZE / raft::warp_size(); + // n_tiles is driven by the *promoted* (4-bit-width) row length, which both operands share at a + // given dim; each side's own native row length only sets how many native bytes + // stage_promoted_tile reads per tile (dim % 32 == 0 guarantees no rounding either way). + const int doc_row_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_document)); + const int query_row_bytes = + static_cast(cuvs::preprocessing::quantize::bbq::get_encoded_row_length(dataset_query)); + const int promoted_row_bytes = static_cast((dataset_query.dim() + 1) / 2); + const int n_tiles = raft::ceildiv(promoted_row_bytes, BBQ_ROW_BYTES); + + const int warp_id_y = warp_id / WARPS_PER_DIM; + const int warp_id_x = warp_id % WARPS_PER_DIM; + + // One phase: accumulate s_row_vec (rows, `new` list) against col_buf over the whole K range, + // then store the accumulators to s_distances (converted to final float distances in place right + // after). col_buf is s_row_vec itself when phase 1 is a + // self-join, otherwise s_col_vec. col_neighbors/col_size select which list the B operand stages. + // alias_tag: staging skip and b_frag source + // row_resident: the A operand is already staged from a previous phase, which is only true + // when the whole row fits one tile (n_tiles == 1) so nothing overwrote it. + auto run_phase = + [&](const Index_t* col_neighbors, int col_size, auto alias_tag, bool row_resident) { + constexpr bool alias_col = decltype(alias_tag)::value; + wmma::fragment c_frag[SUB_PER_DIM][SUB_PER_DIM]; +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + wmma::fill_fragment(c_frag[msub][nsub], 0); + } + } + + for (int step = 0; step < n_tiles; ++step) { + if (!row_resident) { + stage_promoted_tile(s_row_vec, + dataset_document, + new_neighbors, + new_size, + step, + doc_row_bytes, + warp_id, + lane_id); + } + if constexpr (!alias_col) { + stage_promoted_tile(s_col_vec, + dataset_query, + col_neighbors, + col_size, + step, + query_row_bytes, + warp_id, + lane_id); + } + __syncthreads(); + + // a_frag depends only on (msub, kk); b_frag depends only on (nsub, kk) -- load each once + // per kk and reuse across the other sub-tile index, instead of reloading redundantly inside + // a full msub x nsub x kk cross product. + // Deliberately not #pragma unroll'd: full unrolling here keeps more fragment live ranges + // simultaneous, driving register pressure up (64/thread, tied with SMEM for the occupancy + // cap) -- letting the compiler pick reduces that at the cost of some intra-warp ILP. + for (int kk = 0; kk < K_STEPS_PER_TILE; ++kk) { + wmma::fragment + a_frag[SUB_PER_DIM]; + wmma::fragment + b_frag[SUB_PER_DIM]; + const auto* col_buf = alias_col ? s_row_vec : s_col_vec; // compile-time select +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { + const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; + wmma::load_matrix_sync(a_frag[msub], s_row_vec[row0] + kk * (MMA_K / 2), ROW_STRIDE_U4); + } +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; + wmma::load_matrix_sync(b_frag[nsub], col_buf[col0] + kk * (MMA_K / 2), ROW_STRIDE_U4); + } +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + wmma::mma_sync(c_frag[msub][nsub], a_frag[msub], b_frag[nsub], c_frag[msub][nsub]); + } + } + } + __syncthreads(); + } + +#pragma unroll + for (int msub = 0; msub < SUB_PER_DIM; ++msub) { + const int row0 = warp_id_y * WARP_TILE + msub * MMA_M; +#pragma unroll + for (int nsub = 0; nsub < SUB_PER_DIM; ++nsub) { + const int col0 = warp_id_x * WARP_TILE + nsub * MMA_N; + wmma::store_matrix_sync( + reinterpret_cast(s_distances) + row0 * MMA_STORE_STRIDE + col0, + c_frag[msub][nsub], + MMA_STORE_STRIDE, + wmma::mem_row_major); + } + } + __syncthreads(); + + // Converts store_matrix_sync's raw int32 dot products into final float distances, in place + // col = i % MAX_NUM_BI_SAMPLES is invariant across a thread's own iterations + // (BLOCK_SIZE is a multiple of MAX_NUM_BI_SAMPLES), so this thread's column factors are + // fetched once and reused below + const int my_col = tx % MAX_NUM_BI_SAMPLES; + bbq_dequant_factors my_col_factors{}; + if (my_col < col_size) { + my_col_factors = get_dequant_factors(dataset_query, col_neighbors[my_col]); + } + + // Cell (row, col) is read (as raw int, via this alias) and written (as the final float) at + // the same address by the same thread + auto* raw_view = reinterpret_cast(s_distances); + constexpr int total_cells = MAX_NUM_BI_SAMPLES * MAX_NUM_BI_SAMPLES; + for (int i = tx; i < total_cells; i += BLOCK_SIZE) { + const int row = i / MAX_NUM_BI_SAMPLES; + const int col = i % MAX_NUM_BI_SAMPLES; + if (row >= new_size || col >= col_size) continue; + const int distance0 = row * MMA_STORE_STRIDE + col; + const Index_t doc_id = new_neighbors[row]; + const Index_t query_id = col_neighbors[col]; + const uint32_t raw = static_cast(raw_view[distance0]); + s_distances[distance0] = bbq_calculate_metric(raw, + s_document_factors[row], + my_col_factors, + dataset_document, + dataset_query, + metric, + dist_epilogue, + doc_id, + query_id); + } + __syncthreads(); + }; + + // Rows are always new_neighbors/document in both phases below, so this runs exactly once -- see + // the identical staging and rationale in local_join_kernel_bbq_simt. + for (int i = tx; i < new_size; i += BLOCK_SIZE) { + s_document_factors[i] = get_dequant_factors(dataset_document, new_neighbors[i]); + } + __syncthreads(); + + // ---- Phase 1: new x new ---- + run_phase(new_neighbors, new_size, std::integral_constant{}, false); + + for (int step = 0; step < raft::ceildiv(new_size, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + new_neighbors, + s_distances, + true, + MMA_STORE_STRIDE, + new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + if (!old_size) return; + __syncthreads(); + + // ---- Phase 2: new x old ---- + run_phase(old_neighbors, old_size, std::false_type{}, n_tiles == 1); + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= new_size) continue; + auto min_elem = get_min_item(s_list[idx_in_list], + idx_in_list, + old_neighbors, + s_distances, + true, + MMA_STORE_STRIDE, + old_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[idx_in_list], graph, dists, graph_width, locks); + } + } + + for (int step = 0; step < raft::ceildiv(MAX_NUM_BI_SAMPLES, num_warps); ++step) { + const int idx_in_list = step * num_warps + tx / raft::warp_size(); + if (idx_in_list >= old_size) continue; + const int list_idx = idx_in_list + MAX_NUM_BI_SAMPLES; + auto min_elem = get_min_item( + s_list[list_idx], idx_in_list, new_neighbors, s_distances, false, MMA_STORE_STRIDE, new_size); + if (min_elem.id() < gridDim.x) { + insert_to_global_graph(min_elem, s_list[list_idx], graph, dists, graph_width, locks); + } + } +#endif // (__CUDA_ARCH__ >= 750) +} + +// launch_bounds here denote BLOCK_SIZE = 512 and MIN_BLOCKS_PER_SM = 4 +// Per +// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications, +// MAX_RESIDENT_THREAD_PER_SM = BLOCK_SIZE * BLOCKS_PER_SM = 2048 +// For architectures 750 and 860 (890), the values for MAX_RESIDENT_THREAD_PER_SM +// is 1024 and 1536 respectively, which means the bounds don't work anymore +// Used for fp32 data downcast to fp16, and all types using non-L1 distance metric. namespace { template @@ -1147,6 +2055,75 @@ int insert_to_ordered_list(InternalID_t* list, return idx_insert; }; +// Copies `h_graph` (stride `in_degree`) into `out` (stride `out_degree`), dropping self-edges and +// repeated ids, then pads each short list with distinct random ids. +// +// The device-side lists can hold the same id twice: insert_to_global_graph only rejects an +// incoming element when it compares equal to its immediate sorted neighbours, which misses a +// second copy that arrived with a different distance. +template +void shrink_graph_removing_duplicates(Index_t* out, + const InternalID_t* h_graph, + const size_t nrow, + const size_t in_degree, + const size_t out_degree) +{ + // Copy the output graph while removing duplicates. Each thread keeps a bit-packed "seen" + // array, one bit per dataset row, to test and mark ids in O(1) as it scans a row's + // candidates. Only the ids actually placed for a row are ever set, and they're cleared again + // immediately after that row is done, so beyond the one-time zero-initialization when each + // thread starts, no reset across the full array is ever needed. + const size_t num_dedup_words = (nrow + 63) / 64; +#pragma omp parallel + { + std::vector seen_bits(num_dedup_words, 0); + + auto test_and_set = [&](size_t idx) -> bool { + uint64_t mask = uint64_t{1} << (idx & 63); + if (seen_bits[idx >> 6] & mask) { return false; } + seen_bits[idx >> 6] |= mask; + return true; + }; + +#pragma omp for + for (size_t i = 0; i < nrow; i++) { + auto* output_neighbor_list_ptr = out + i * out_degree; + + size_t out_j = 0; + // Copy neighbor list while removing duplicates. + for (size_t in_j = 0; in_j < out_degree; in_j++) { + size_t idx = h_graph[i * in_degree + in_j].id(); + // Unfilled slots carry a sentinel id; leave them to the random fill below. + if (idx >= nrow || idx == i || !test_and_set(idx)) { continue; } + output_neighbor_list_ptr[out_j] = static_cast(idx); + out_j++; + } + + // Fill with random nodes if the length of the filled neighbor list is less than the degree. + for (size_t j = out_j; j < out_degree; j++) { + uint64_t rnd = static_cast(i * out_degree + j + 1); + uint64_t idx = 0; + bool placed = false; + for (size_t attempts = 0; !placed && attempts < out_degree; attempts++) { + rnd = cuvs::neighbors::detail::device::xorshift64(rnd); + idx = rnd % nrow; + placed = (idx != i) && test_and_set(idx); + } + output_neighbor_list_ptr[j] = static_cast(idx); + } + + // Unset every bit this row touched so the array is back to all-zero for the next row this + // thread processes. Since seen_bits is all-zero on entry to this row and thread-local, the + // only bits set anywhere are ones this row's own entries set, so the whole word covering + // idx can be zeroed outright rather than masking off just its one bit. + for (size_t k = 0; k < out_degree; k++) { + size_t idx = static_cast(output_neighbor_list_ptr[k]); + seen_bits[idx >> 6] = 0; + } + } + } +} + } // namespace template @@ -1496,6 +2473,145 @@ void GNND::local_join(cudaStream_t stream, DistEpilogue_t dist_ } } +template +template +void GNND::local_join( + cudaStream_t stream, + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + DistEpilogue_t dist_epilogue) +{ + raft::matrix::fill(res, dists_buffer_.view(), std::numeric_limits::max()); + + // Both kernels take the same (document, query) pair, so there is no symmetric/asymmetric split + // here: a single quantizer just means the same one on both operands, which is exactly what + // SelfJoin encodes. Picking the two quantizers is all that differs. + const bool self_join = dataset.quantizers.size() == 1; + const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + + // Asymmetric: a packed_4b query selects the tensor-core path, a transposed query the SIMT one. + // Only packed_1b promotes to the tensor-core path; transposed_2b is SIMT-only (it would need + // promoting from two bitplanes at once, which the int4 MMA staging doesn't support). + const bool tc_pair = has_4b && has_1b; + const bool simt_pair = (has_4bt && (has_1b || has_2bt)) || (has_2bt && has_1b); + RAFT_EXPECTS(self_join || tc_pair || simt_pair, + "Unsupported BBQ layout pair for asymmetric local join. Supported: " + "packed_1b x packed_4b (tensor core); packed_1b x transposed_2b, " + "packed_1b x transposed_4b, transposed_2b x transposed_4b (SIMT)."); + auto quantizer_query = + self_join ? dataset.quantizers[0] + : (tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) + : (has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b))); + auto quantizer_document = self_join ? dataset.quantizers[0] + : has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) + : dataset.get_quantizer(bbq_code_layout::transposed_2b); + + // stage_tile_simt / stage_promoted_tile cast code buffers to uint32_t*, so every plane stride + // must be 4-byte aligned. + { + const auto len = cuvs::preprocessing::quantize::bbq::get_encoded_row_length(quantizer_query); + const int n_planes = + cuvs::preprocessing::quantize::bbq::get_code_planes(quantizer_query.layout); + RAFT_EXPECTS(len % (4u * static_cast(n_planes)) == 0, + "BBQ local join requires the encoded row length to be a multiple of 4*n_planes " + "for 32-bit aligned plane loads, got %u with n_planes = %d", + len, + n_planes); + RAFT_EXPECTS(quantizer_document.dim() % 32 == 0, + "BBQ local join requires dataset dim to be a multiple of 32, got %lld", + static_cast(quantizer_document.dim())); + } + + // One launch site for both kernels: they take identical arguments, and the query's layout picks + // the path -- packed_4b is the only layout the int4 tensor-core kernel is dispatched for. + auto launch = [&](auto document_layout, auto query_layout, auto self_join_tag) { + constexpr auto D = decltype(document_layout)::value; + constexpr auto Q = decltype(query_layout)::value; + constexpr bool S = decltype(self_join_tag)::value; + auto launch_local_join = [&](auto kernel) { + kernel<<>>(graph_.h_graph_new.data_handle(), + h_rev_graph_new_.data_handle(), + d_list_sizes_new_.data_handle(), + h_graph_old_.data_handle(), + h_rev_graph_old_.data_handle(), + d_list_sizes_old_.data_handle(), + NUM_SAMPLES, + quantizer_document, + quantizer_query, + graph_buffer_.data_handle(), + dists_buffer_.data_handle(), + DEGREE_ON_DEVICE, + d_locks_.data_handle(), + build_config_.metric, + dist_epilogue); + }; + // Naming a kernel rather than launching it directly leaves the trailing template parameters + // nothing to deduce from, so they are spelled out here. + using kernel_data_t = std::remove_const_t; + using kernel_id_t = InternalID_t; + if constexpr (Q == bbq_code_layout::packed_4b) { + launch_local_join( + local_join_kernel_bbq_wmma); + } else { + launch_local_join( + local_join_kernel_bbq_simt); + } + }; + const bbq_code_layout d = quantizer_document.layout; + const bbq_code_layout q = quantizer_query.layout; + if (self_join) { + switch (d) { + case bbq_code_layout::packed_1b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::transposed_2b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_4b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_7b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + case bbq_code_layout::packed_8b: + launch(std::integral_constant{}, + std::integral_constant{}, + std::true_type{}); + break; + default: RAFT_FAIL("Unsupported BBQ layout for symmetric local join on this branch."); + } + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::packed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::transposed_2b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::packed_1b && q == bbq_code_layout::transposed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else if (d == bbq_code_layout::transposed_2b && q == bbq_code_layout::transposed_4b) { + launch(std::integral_constant{}, + std::integral_constant{}, + std::false_type{}); + } else { + RAFT_FAIL("Unsupported BBQ layout pair for asymmetric local join."); + } +} + template template void GNND::build(Data_t* data, @@ -1727,71 +2843,216 @@ void GNND::build(Data_t* data, Index_t* graph_shrink_buffer = (Index_t*)graph_.h_dists.data_handle(); - // Copy the output graph while removing duplicates. Each thread keeps a bit-packed "seen" - // array, one bit per dataset row, to test and mark ids in O(1) as it scans a row's - // candidates. Only the ids actually placed for a row are ever set, and they're cleared again - // immediately after that row is done, so beyond the one-time zero-initialization when each - // thread starts, no reset across the full array is ever needed. - const size_t num_dedup_words = (static_cast(nrow_) + 63) / 64; -#pragma omp parallel - { - std::vector seen_bits(num_dedup_words, 0); + shrink_graph_removing_duplicates( + graph_shrink_buffer, graph_.h_graph, nrow_, graph_.node_degree, build_config_.node_degree); + graph_.h_graph = nullptr; - auto test_and_set = [&](size_t idx) -> bool { - uint64_t mask = uint64_t{1} << (idx & 63); - if (seen_bits[idx >> 6] & mask) { return false; } - seen_bits[idx >> 6] |= mask; - return true; - }; +#pragma omp parallel for + for (size_t i = 0; i < (size_t)nrow_; i++) { + for (size_t j = 0; j < build_config_.node_degree; j++) { + output_graph[i * build_config_.node_degree + j] = + graph_shrink_buffer[i * build_config_.node_degree + j]; + } + } +} -#pragma omp for - for (size_t i = 0; i < (size_t)nrow_; i++) { - auto output_neighbor_list_ptr = graph_shrink_buffer + i * build_config_.node_degree; +template +template +void GNND::build( + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + Index_t* output_graph, + bool return_distances, + DistData_t* output_distances, + DistEpilogue_t dist_epilogue) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); + nrow_ = static_cast(dataset.n_rows()); + graph_.nrow = nrow_; + graph_.bloom_filter.set_nrow(nrow_); + update_counter_ = 0; + graph_.h_graph = reinterpret_cast*>(output_graph); - size_t out_j = 0; + graph_.clear(); + graph_.init_random_graph(); + graph_.sample_graph(true); - // Copy neighbor list while removing duplicates. - for (size_t in_j = 0; in_j < build_config_.node_degree; in_j++) { - size_t idx = graph_.h_graph[i * graph_.node_degree + in_j].id(); - if (idx >= (size_t)nrow_ || idx == i || !test_and_set(idx)) continue; - output_neighbor_list_ptr[out_j] = idx; - out_j++; + auto update_and_sample = [&](bool update_graph) { + if (update_graph) { + update_counter_ = 0; + graph_.update_graph(graph_host_buffer_.data_handle(), + dists_host_buffer_.data_handle(), + DEGREE_ON_DEVICE, + update_counter_); + if (update_counter_ < build_config_.termination_threshold * nrow_ * + build_config_.dataset_dim / counter_interval) { + update_counter_ = -1; } + } + graph_.sample_graph(false); + }; - // Fill with random nodes if the length of the filled neighbor list is less than the degree. - for (size_t j = out_j; j < build_config_.node_degree; j++) { - uint64_t rnd = static_cast(i * build_config_.node_degree + j + 1); - uint64_t idx = 0; - bool placed = false; - for (size_t attempts = 0; !placed && attempts < build_config_.node_degree; attempts++) { - rnd = cuvs::neighbors::detail::device::xorshift64(rnd); - idx = rnd % nrow_; - placed = (idx != i) && test_and_set(idx); - } - output_neighbor_list_ptr[j] = static_cast(idx); - } + for (size_t it = 0; it < build_config_.max_iterations; ++it) { + raft::copy(res, d_list_sizes_new_.view(), graph_.h_list_sizes_new.view()); + raft::copy(res, h_graph_old_.view(), graph_.h_graph_old.view()); + raft::copy(res, d_list_sizes_old_.view(), graph_.h_list_sizes_old.view()); + raft::resource::sync_stream(res); - // Unset every bit this row touched so the array is back to all-zero for the next row this - // thread processes. Since seen_bits is all-zero on entry to this row and thread-local, the - // only bits set anywhere are ones this row's own entries set, so the whole word covering - // idx can be zeroed outright rather than masking off just its one bit. - for (size_t k = 0; k < build_config_.node_degree; k++) { - size_t idx = static_cast(output_neighbor_list_ptr[k]); - seen_bits[idx >> 6] = 0; + std::thread update_and_sample_thread(update_and_sample, it); + RAFT_LOG_DEBUG("# GNND iteration: %lu / %lu", it + 1, build_config_.max_iterations); + + static_assert(DEGREE_ON_DEVICE * sizeof(*(dists_buffer_.data_handle())) >= + NUM_SAMPLES * sizeof(*(graph_buffer_.data_handle()))); + add_reverse_edges(graph_.h_graph_new.data_handle(), + h_rev_graph_new_.data_handle(), + reinterpret_cast(dists_buffer_.data_handle()), + d_list_sizes_new_.data_handle(), + stream); + add_reverse_edges(h_graph_old_.data_handle(), + h_rev_graph_old_.data_handle(), + reinterpret_cast(dists_buffer_.data_handle()), + d_list_sizes_old_.data_handle(), + stream); + + local_join(stream, dataset, dist_epilogue); + update_and_sample_thread.join(); + if (update_counter_ == -1) { break; } + raft::copy(res, graph_host_buffer_.view(), graph_buffer_.view()); + raft::copy(res, dists_host_buffer_.view(), dists_buffer_.view()); + raft::resource::sync_stream(res); + graph_.sample_graph_new(graph_host_buffer_.data_handle(), DEGREE_ON_DEVICE); + } + + graph_.update_graph(graph_host_buffer_.data_handle(), + dists_host_buffer_.data_handle(), + DEGREE_ON_DEVICE, + update_counter_); + raft::resource::sync_stream(res); + graph_.sort_lists(); + + static_assert(sizeof(decltype(*(graph_.h_dists.data_handle()))) >= sizeof(Index_t)); + if (return_distances) { + auto graph_h_dists = raft::make_host_matrix( + nrow_, build_config_.output_graph_degree); +#pragma omp parallel for + for (size_t i = 0; i < nrow_; ++i) { + for (size_t j = 0; j < build_config_.output_graph_degree; ++j) { + graph_h_dists(i, j) = graph_.h_dists(i, j); } } + raft::copy( + res, + raft::make_device_vector_view(output_distances, nrow_ * build_config_.output_graph_degree), + raft::make_host_vector_view(graph_h_dists.data_handle(), + nrow_ * build_config_.output_graph_degree)); + + auto output_dist_view = raft::make_device_matrix_view( + output_distances, nrow_, build_config_.output_graph_degree); + const bool can_postprocess_dist = std::is_same_v; + if (build_config_.metric == cuvs::distance::DistanceType::L2SqrtExpanded && + can_postprocess_dist) { + raft::linalg::map( + res, output_dist_view, raft::sqrt_op{}, raft::make_const_mdspan(output_dist_view)); + } else if (!cuvs::distance::is_min_close(build_config_.metric) && can_postprocess_dist) { + raft::linalg::map(res, + output_dist_view, + raft::mul_const_op(-1), + raft::make_const_mdspan(output_dist_view)); + } + raft::resource::sync_stream(res); } + + auto* graph_shrink_buffer = reinterpret_cast(graph_.h_dists.data_handle()); + + shrink_graph_removing_duplicates( + graph_shrink_buffer, graph_.h_graph, nrow_, graph_.node_degree, build_config_.node_degree); graph_.h_graph = nullptr; #pragma omp parallel for - for (size_t i = 0; i < (size_t)nrow_; i++) { - for (size_t j = 0; j < build_config_.node_degree; j++) { + for (size_t i = 0; i < nrow_; ++i) { + for (size_t j = 0; j < build_config_.node_degree; ++j) { output_graph[i * build_config_.node_degree + j] = graph_shrink_buffer[i * build_config_.node_degree + j]; } } } +template +void build(raft::resources const& res, + const index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + index& idx) +{ + RAFT_EXPECTS(dataset.quantizers.size() > 0, "BBQ dataset must not be empty."); + auto front_quantizer = dataset.quantizers[0]; + cuvs::common::nvtx::range fun_scope( + "neighbors::nn_descent::detail::build-bbq(%zu, %zu, %zu, %zu, %zu)", + size_t(dataset.n_rows()), + size_t(dataset.dim()), + size_t(idx.graph().extent(1)), + size_t(idx.metric()), + size_t(cuvs::preprocessing::quantize::bbq::get_bit_width(front_quantizer.layout))); + RAFT_EXPECTS(idx.metric() == cuvs::distance::DistanceType::L2Expanded || + idx.metric() == cuvs::distance::DistanceType::L2SqrtExpanded || + idx.metric() == cuvs::distance::DistanceType::CosineExpanded || + idx.metric() == cuvs::distance::DistanceType::InnerProduct, + "BBQ NN-Descent supports L2Expanded, L2SqrtExpanded, CosineExpanded, and " + "InnerProduct."); + RAFT_EXPECTS(idx.metric() == front_quantizer.metric, + "BBQ dataset metric does not match the NN-Descent metric."); + // packed_4b is the only layout dispatched to local_join_kernel_bbq_wmma, and the int4 MMA that + // kernel is built around only exists from sm_75 on. + if (dataset.has_layout(bbq_code_layout::packed_4b)) { + auto kernel = local_join_kernel_bbq_wmma, + raft::identity_op>; + auto runtime_arch = raft::util::arch::kernel_virtual_arch(reinterpret_cast(kernel)); + RAFT_EXPECTS( + raft::util::arch::SM_range(raft::util::arch::SM_75(), raft::util::arch::SM_future()) + .contains(runtime_arch), + "The BBQ packed_4b layout requires int4 tensor cores (compute capability 7.5 or newer), but " + "the local join kernel resolves to %d.%d here. Use transposed_4b for 4-bit codes instead.", + runtime_arch.value() / 100, + (runtime_arch.value() / 10) % 10); + } + + size_t extended_graph_degree; + size_t graph_degree; + auto build_config = get_build_config(res, + params, + dataset.n_rows(), + dataset.dim(), + idx.metric(), + extended_graph_degree, + graph_degree); + auto int_graph = + raft::make_host_matrix(dataset.n_rows(), extended_graph_degree); + GNND nnd(res, build_config); + + if (idx.distances().has_value() || !params.return_distances) { + nnd.build(dataset, + int_graph.data_handle(), + params.return_distances, + idx.distances() + .value_or(raft::make_device_matrix(res, 0, 0).view()) + .data_handle()); + } else { + RAFT_FAIL( + "Distance view not allocated. Using return_distances set to true requires " + "distance view to be allocated."); + } + +#pragma omp parallel for + for (size_t i = 0; i < static_cast(dataset.n_rows()); ++i) { + for (size_t j = 0; j < graph_degree; ++j) { + idx.graph()(i, j) = int_graph(i, j); + } + } +} + template +index build(raft::resources const& res, + const index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset) +{ + size_t graph_degree = params.graph_degree; + if (params.intermediate_graph_degree < graph_degree) { + RAFT_LOG_WARN( + "Graph degree (%lu) cannot be larger than intermediate graph degree (%lu), reducing " + "graph_degree.", + graph_degree, + params.intermediate_graph_degree); + graph_degree = params.intermediate_graph_degree; + } + + index idx{res, + static_cast(dataset.n_rows()), + static_cast(graph_degree), + params.return_distances, + params.metric}; + detail::build(res, params, dataset, idx); + return idx; +} + template + void build(cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + Index_t* output_graph, + bool return_distances, + DistData_t* output_distances, + DistEpilogue_t dist_epilogue = DistEpilogue_t{}); ~GNND() = default; using ID_t = InternalID_t; void reset(raft::resources const& res); @@ -219,6 +226,12 @@ class CUVS_EXPORT GNND { template void local_join(cudaStream_t stream = 0, DistEpilogue_t dist_epilogue = DistEpilogue_t{}); + template + void local_join( + cudaStream_t stream, + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + DistEpilogue_t dist_epilogue = DistEpilogue_t{}); + raft::resources const& res; BuildConfig build_config_; diff --git a/cpp/src/neighbors/nn_descent.cuh b/cpp/src/neighbors/nn_descent.cuh index abe45f0a6f..9c9b3c5540 100644 --- a/cpp/src/neighbors/nn_descent.cuh +++ b/cpp/src/neighbors/nn_descent.cuh @@ -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 */ @@ -184,6 +184,23 @@ void build(raft::resources const& res, detail::build(res, params, dataset, idx); } +template +auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset) -> index +{ + return detail::build(res, params, dataset); +} + +template +void build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + index& idx) +{ + detail::build(res, params, dataset, idx); +} + /** @} */ // end group nn-descent } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/src/neighbors/nn_descent_gnnd_inst.cu b/cpp/src/neighbors/nn_descent_gnnd_inst.cu index 0e4bb3a534..a5b2517b9c 100644 --- a/cpp/src/neighbors/nn_descent_gnnd_inst.cu +++ b/cpp/src/neighbors/nn_descent_gnnd_inst.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -39,6 +39,12 @@ template void detail::GNND::build( bool return_distances, float* output_distances, raft::identity_op dist_epilogue); +template void detail::GNND::build( + cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + int* output_graph, + bool return_distances, + float* output_distances, + raft::identity_op dist_epilogue); template void detail::GNND::local_join( cudaStream_t stream, raft::identity_op dist_epilogue); diff --git a/cpp/src/neighbors/nn_descent_inst.cu.in b/cpp/src/neighbors/nn_descent_inst.cu.in index fbbeb58f7a..62d749e5b6 100644 --- a/cpp/src/neighbors/nn_descent_inst.cu.in +++ b/cpp/src/neighbors/nn_descent_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -52,4 +52,21 @@ auto build(raft::resources const& handle, } }; +auto build(raft::resources const& handle, + const cuvs::neighbors::nn_descent::index_params& params, + cuvs::neighbors::device_bbq_dataset_view dataset, + std::optional> graph) + -> cuvs::neighbors::nn_descent::index +{ + if (!graph.has_value()) { + return cuvs::neighbors::nn_descent::build(handle, params, dataset); + } else { + std::optional> distances = + std::nullopt; + cuvs::neighbors::nn_descent::index idx{ + handle, graph.value(), distances, params.metric}; + cuvs::neighbors::nn_descent::build(handle, params, dataset, idx); + return idx; + }; +} } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/src/preprocessing/quantize/bbq.cu b/cpp/src/preprocessing/quantize/bbq.cu new file mode 100644 index 0000000000..991f47b77e --- /dev/null +++ b/cpp/src/preprocessing/quantize/bbq.cu @@ -0,0 +1,48 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +namespace preprocessing::quantize::bbq { + +namespace helpers { +void resolve_dequant_factors( + raft::resources const& res, + raft::device_vector_view dequant_delta, + raft::device_vector_view dequant_sum_delta, + raft::device_vector_view lower_intervals, + raft::device_vector_view upper_intervals, + raft::device_vector_view quantized_component_sums, + bbq_code_layout layout) +{ + const auto n_rows = dequant_delta.extent(0); + RAFT_EXPECTS(dequant_sum_delta.extent(0) == n_rows && lower_intervals.extent(0) == n_rows && + upper_intervals.extent(0) == n_rows && + quantized_component_sums.extent(0) == n_rows, + "resolve_dequant_factors: all vectors must have the same length"); + const float scale = 1.0f / static_cast((uint32_t{1} << get_bit_width(layout)) - 1); + auto* sum_delta = dequant_sum_delta.data_handle(); + raft::linalg::map_offset(res, + dequant_delta, + [scale, + sum_delta, + lower = lower_intervals.data_handle(), + upper = upper_intervals.data_handle(), + sums = quantized_component_sums.data_handle()] __device__(int64_t i) { + const float delta = (upper[i] - lower[i]) * scale; + sum_delta[i] = delta * static_cast(sums[i]); + return delta; + }); +} +} // namespace helpers +} // namespace preprocessing::quantize::bbq +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh new file mode 100644 index 0000000000..4d426d5454 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/bbq_distance.cuh @@ -0,0 +1,515 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace preprocessing::quantize::bbq { + +#ifdef __CUDACC__ + +// -------------------------------------------------------------------------- +// Layout geometry +// Encoded row length of a quantizer, and the bit planes its layout slices a row into. +// -------------------------------------------------------------------------- + +template +_RAFT_HOST_DEVICE constexpr uint32_t get_encoded_row_length( + const quantizer_view& dataset) +{ + return get_encoded_row_length(dataset.dim(), dataset.layout); +} + +/** + * Bit planes a row is sliced into: the transposed layouts hold one plane per code bit, the packed + * ones are a single dense plane. This is all the tiling and the cross-plane inner products need + * to know about a layout. + */ +_RAFT_HOST_DEVICE constexpr int get_code_planes(const bbq_code_layout layout) +{ + switch (layout) { + case bbq_code_layout::transposed_2b: return 2; + case bbq_code_layout::transposed_4b: return 4; + default: return 1; + } +} + +// -------------------------------------------------------------------------- +// Code inner products (1x1) +// Raw byte-level dot products over one row pair, per code layout. +// -------------------------------------------------------------------------- + +/** + * Cross-plane binary inner product over `Planes` bit planes, shifting each (i, j) plane pair by + * i + j. `Planes == 1` is a plain binary product, so this covers packed_1b as well as the + * transposed layouts. The 4-byte body needs both rows 4-byte aligned; the byte tail handles a + * stripe whose length is not a multiple of 4. + */ +template +__device__ __forceinline__ uint32_t code_inner_product_transposed(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t result = 0) +{ + const size_t stripe = n_bytes / Planes; + // Plane p starts at row + p * stripe and is read as uint32_t words, so a stripe that is not a + // multiple of 4 misaligns every plane past the first -- the byte tail below only covers a ragged + // stripe *length*, not a ragged stripe *offset*. Planes == 1 has no offset and so is exempt. + assert(Planes == 1 || stripe % sizeof(uint32_t) == 0); +#pragma unroll + for (int i = 0; i < Planes; ++i) { +#pragma unroll + for (int j = 0; j < Planes; ++j) { + const uint8_t* a = row_a + i * stripe; + const uint8_t* b = row_b + j * stripe; + uint32_t partial = 0; + size_t k = 0; +#pragma unroll 4 + for (; k + sizeof(uint32_t) <= stripe; k += sizeof(uint32_t)) { + partial += __popc(*reinterpret_cast(a + k) & + *reinterpret_cast(b + k)); + } + for (; k < stripe; ++k) { + partial += __popc(static_cast(a[k] & b[k])); + } + result += partial << (i + j); + } + } + return result; +} + +/** One word of two packed_4b rows: the two nibble halves are two masked dp4a products. */ +__device__ __forceinline__ uint32_t dp4a_packed_4b_word(uint32_t a, uint32_t b, uint32_t total) +{ + constexpr uint32_t nibble_mask = 0x0F0F0F0Fu; + total = __dp4a(a & nibble_mask, b & nibble_mask, total); + return __dp4a((a >> 4) & nibble_mask, (b >> 4) & nibble_mask, total); +} + +/** Symmetric for packNibbles (Lucene int4DotProductBothPacked). */ +__device__ __forceinline__ uint32_t code_inner_product_packed_4b(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t total = 0) +{ + size_t i = 0; +#pragma unroll 4 + for (; i + 4 <= n_bytes; i += 4) { + total = dp4a_packed_4b_word(*reinterpret_cast(row_a + i), + *reinterpret_cast(row_b + i), + total); + } + for (; i < n_bytes; ++i) { + const unsigned a = row_a[i]; + const unsigned b = row_b[i]; + total += (a & 0x0Fu) * (b & 0x0Fu); + total += ((a >> 4) & 0x0Fu) * ((b >> 4) & 0x0Fu); + } + return total; +} + +/** One-byte-per-code dot product, optionally masking unused high bits. */ +__device__ __forceinline__ uint32_t code_inner_product_packed_8b(const uint8_t* row_a, + const uint8_t* row_b, + size_t n_bytes, + uint32_t result = 0, + uint8_t code_mask = 0xFFu) +{ + const uint32_t word_mask = uint32_t{code_mask} * 0x01010101u; + size_t i = 0; +#pragma unroll 4 + for (; i + 4 <= n_bytes; i += 4) { + const auto a = *reinterpret_cast(row_a + i) & word_mask; + const auto b = *reinterpret_cast(row_b + i) & word_mask; + result = __dp4a(a, b, result); + } + for (; i < n_bytes; ++i) { + result += + static_cast(row_a[i] & code_mask) * static_cast(row_b[i] & code_mask); + } + return result; +} + +/** + * Integer inner product between two encoded rows. + * + * The uint32_t result bounds every BBQ layout to 66,050 dimensions: the worst case is + * `packed_8b`, where `dim * 255 * 255` must not exceed UINT32_MAX. + */ +__device__ __forceinline__ uint32_t code_inner_product(const uint8_t* row_a, + const uint8_t* row_b, + const bbq_code_layout layout, + const size_t n_bytes, + uint32_t result = 0) +{ + switch (layout) { + case bbq_code_layout::packed_1b: + return code_inner_product_transposed<1>(row_a, row_b, n_bytes, result); + case bbq_code_layout::transposed_2b: + return code_inner_product_transposed<2>(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_4b: + return code_inner_product_packed_4b(row_a, row_b, n_bytes, result); + case bbq_code_layout::transposed_4b: + return code_inner_product_transposed<4>(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_8b: + return code_inner_product_packed_8b(row_a, row_b, n_bytes, result); + case bbq_code_layout::packed_7b: + default: + return code_inner_product_packed_8b( + row_a, + row_b, + n_bytes, + result, + static_cast((uint32_t{1} << get_bit_width(layout)) - 1)); + } +} + +// -------------------------------------------------------------------------- +// Code promotion +// Widens a narrower layout to 4-bit width, so an asymmetric pair can meet in one format. +// -------------------------------------------------------------------------- + +// Promotes one native word of dense packed_1b codes (1 bit/value, 8 values/byte, MSB-first: +// the value at position 8*byte+i sits at bit (7-i)) into four 4-bit-width, packed_4b-style +// output words, packed into one uint4 for a single vectorized store. +// +// Branch-free SWAR: extract the 4 (2-bit) fields of each byte lane-wise across all 4 native +// bytes at once, spread each field 0-3 into a nibble value lane-wise, then transpose the 4 +// resulting field-words into the 4 per-native-byte output words with chained __byte_perm pairs. +__device__ __forceinline__ uint4 packed_1b_to_4b(uint32_t native_word) +{ + const uint32_t spread0 = ((native_word >> 3) & 0x10101010u) | ((native_word >> 6) & 0x01010101u); + const uint32_t spread1 = ((native_word >> 1) & 0x10101010u) | ((native_word >> 4) & 0x01010101u); + const uint32_t spread2 = ((native_word << 1) & 0x10101010u) | ((native_word >> 2) & 0x01010101u); + const uint32_t spread3 = ((native_word << 3) & 0x10101010u) | (native_word & 0x01010101u); + const uint32_t t0 = __byte_perm(spread0, spread1, 0x5140); + const uint32_t t1 = __byte_perm(spread0, spread1, 0x7362); + const uint32_t t2 = __byte_perm(spread2, spread3, 0x5140); + const uint32_t t3 = __byte_perm(spread2, spread3, 0x7362); + return uint4{__byte_perm(t0, t2, 0x5410), + __byte_perm(t0, t2, 0x7632), + __byte_perm(t1, t3, 0x5410), + __byte_perm(t1, t3, 0x7632)}; +} + +// -------------------------------------------------------------------------- +// Cross-layout inner products (1x1) +// One document row against one query row encoded in a different layout. +// -------------------------------------------------------------------------- + +/** + * Cross-plane binary inner product over two bit-sliced rows whose plane counts need not match, + * shifting each (i, j) plane pair by i + j. The runtime-width counterpart of + * code_inner_product_transposed, which the local-join kernels reach with both counts known at + * compile time; a single plane is a plain binary product, so this covers packed_1b too. + * + * All three bit-sliced layouts store ceildiv(dim, 8) bytes per plane, which is what lets one + * @p plane_bytes describe both operands. + */ +__device__ __forceinline__ uint32_t code_inner_product_planes(const uint8_t* row_document, + int document_planes, + const uint8_t* row_query, + int query_planes, + size_t plane_bytes, + uint32_t result = 0) +{ + // Planes past the first start at a multiple of plane_bytes and are read as uint32_t words. + assert((document_planes == 1 && query_planes == 1) || plane_bytes % sizeof(uint32_t) == 0); + for (int i = 0; i < document_planes; ++i) { + for (int j = 0; j < query_planes; ++j) { + const uint8_t* a = row_document + i * plane_bytes; + const uint8_t* b = row_query + j * plane_bytes; + uint32_t partial = 0; + size_t k = 0; +#pragma unroll 4 + for (; k + sizeof(uint32_t) <= plane_bytes; k += sizeof(uint32_t)) { + partial += __popc(*reinterpret_cast(a + k) & + *reinterpret_cast(b + k)); + } + for (; k < plane_bytes; ++k) { + partial += __popc(static_cast(a[k] & b[k])); + } + result += partial << (i + j); + } + } + return result; +} + +/** + * packed_1b document against a packed_4b query. The document is promoted to 4-bit width one + * native word at a time -- 32 dimensions, which is exactly the four query words covering the same + * range -- and the pair is multiplied as two packed_4b rows. The SIMT equivalent of what + * stage_promoted_tile plus the u4 MMA do for this layout pair in the tensor-core local join, and + * it pairs codes the same way, since packed_1b_to_4b emits packed_4b-style words. + * + * Requires dim % 32 == 0, so that the promoted document covers the query row exactly. + */ +__device__ __forceinline__ uint32_t code_inner_product_1b_x_packed_4b(const uint8_t* row_document, + const uint8_t* row_query, + size_t document_bytes, + uint32_t result = 0) +{ + assert(document_bytes % sizeof(uint32_t) == 0); + for (size_t i = 0; i + sizeof(uint32_t) <= document_bytes; i += sizeof(uint32_t)) { + uint32_t promoted[4]; + reinterpret_cast(promoted) = + packed_1b_to_4b(*reinterpret_cast(row_document + i)); + const auto* query_words = reinterpret_cast(row_query + 4 * i); +#pragma unroll + for (int e = 0; e < 4; ++e) { + result = dp4a_packed_4b_word(promoted[e], query_words[e], result); + } + } + return result; +} + +/** + * Integer inner product between a document row and a query row, each read through its own + * quantizer. The two layouts may differ: the supported (document, query) pairs are the ones + * nn-descent's local join dispatches on, and passing the same view twice is the symmetric case. + */ +template +__device__ __forceinline__ uint32_t +code_inner_product(const quantizer_view& quantizer_document, + const quantizer_view& quantizer_query, + int64_t row_document, + int64_t row_query) +{ + const uint8_t* document = &quantizer_document.codes(row_document, 0); + const uint8_t* query = &quantizer_query.codes(row_query, 0); + + if (quantizer_document.layout == quantizer_query.layout) { + return code_inner_product( + document, query, quantizer_document.layout, get_encoded_row_length(quantizer_document)); + } + if (quantizer_document.layout == bbq_code_layout::packed_1b && + quantizer_query.layout == bbq_code_layout::packed_4b) { + return code_inner_product_1b_x_packed_4b( + document, query, get_encoded_row_length(quantizer_document)); + } + // Every remaining supported pair is bit-sliced on both sides: (packed_1b, transposed_2b), + // (packed_1b, transposed_4b) and (transposed_2b, transposed_4b). + assert(quantizer_document.layout == bbq_code_layout::packed_1b || + quantizer_document.layout == bbq_code_layout::transposed_2b); + assert(quantizer_query.layout == bbq_code_layout::transposed_2b || + quantizer_query.layout == bbq_code_layout::transposed_4b); + return code_inner_product_planes(document, + get_code_planes(quantizer_document.layout), + query, + get_code_planes(quantizer_query.layout), + (quantizer_document.dim() + 7) / 8); +} + +// -------------------------------------------------------------------------- +// Fused inner products (2x1) +// Two left rows against a shared right operand, for the local-join inner loop. +// -------------------------------------------------------------------------- + +/** + * Two cross-plane inner products over statically known document and query plane counts. + * Not asymmetric-specific: a symmetric pair is just document_planes == query_planes, which is how + * the 1x1 and 2t x 2t self-joins are computed. + */ +template +__device__ inline void code_inner_product_planes_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + constexpr size_t document_plane_stride = document_row_bytes / document_planes; + constexpr size_t query_plane_stride = query_row_bytes / query_planes; + // Both operands are stepped by their own plane stride and then read as uint32_t, so both + // strides -- not just the query's -- must be 4-byte aligned. + static_assert(query_plane_stride % sizeof(uint32_t) == 0); + static_assert(document_plane_stride % sizeof(uint32_t) == 0); +#pragma unroll + for (int p_query = 0; p_query < query_planes; ++p_query) { +#pragma unroll + for (int p_document = 0; p_document < document_planes; ++p_document) { + const uint8_t* a0 = row_a0 + p_document * document_plane_stride; + const uint8_t* a1 = row_a1 + p_document * document_plane_stride; + const uint8_t* b = row_b + p_query * query_plane_stride; + uint32_t partial0 = 0; + uint32_t partial1 = 0; +#pragma unroll 4 + for (size_t i = 0; i < query_plane_stride; i += sizeof(uint32_t)) { + const auto wa0 = *reinterpret_cast(a0 + i); + const auto wa1 = *reinterpret_cast(a1 + i); + const auto wb = *reinterpret_cast(b + i); + partial0 += __popc(wa0 & wb); + partial1 += __popc(wa1 & wb); + } + total0 += partial0 << (p_document + p_query); + total1 += partial1 << (p_document + p_query); + } + } +} + +template +__device__ inline void code_inner_product_packed_4b_symmetric_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + static_assert(n_bytes % sizeof(uint32_t) == 0); + constexpr uint32_t nibble_mask = 0x0F0F0F0Fu; +#pragma unroll 4 + for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { + const auto a0 = *reinterpret_cast(row_a0 + i); + const auto a1 = *reinterpret_cast(row_a1 + i); + const auto b = *reinterpret_cast(row_b + i); + const auto b_low = b & nibble_mask; + const auto b_high = (b >> 4) & nibble_mask; + total0 = __dp4a(a0 & nibble_mask, b_low, total0); + total0 = __dp4a((a0 >> 4) & nibble_mask, b_high, total0); + total1 = __dp4a(a1 & nibble_mask, b_low, total1); + total1 = __dp4a((a1 >> 4) & nibble_mask, b_high, total1); + } +} + +template +__device__ inline void code_inner_product_packed_8b_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1, + uint8_t code_mask = 0xFFu) +{ + static_assert(n_bytes % sizeof(uint32_t) == 0); + const uint32_t word_mask = uint32_t{code_mask} * 0x01010101u; +#pragma unroll 4 + for (size_t i = 0; i < n_bytes; i += sizeof(uint32_t)) { + const auto a0 = *reinterpret_cast(row_a0 + i) & word_mask; + const auto a1 = *reinterpret_cast(row_a1 + i) & word_mask; + const auto b = *reinterpret_cast(row_b + i) & word_mask; + total0 = __dp4a(a0, b, total0); + total1 = __dp4a(a1, b, total1); + } +} + +// Selects the SIMT inner product for a (document, query) layout pair: bit-sliced layouts go to the +// cross-plane popc, densely-packed ones to dp4a. dp4a needs both operands in the same packing, so +// those forms apply whenever the two layouts match. packed_7b and packed_8b reach this from +// GNND::local_join; packed_4b does not (symmetric packed_4b goes to the wmma kernel), but is kept +// as a SIMT reference point. +template +__device__ __forceinline__ void bbq_code_inner_product_2x1(const uint8_t* row_a0, + const uint8_t* row_a1, + const uint8_t* row_b, + uint32_t& total0, + uint32_t& total1) +{ + namespace bbq = cuvs::preprocessing::quantize::bbq; + if constexpr (DocumentLayout == QueryLayout && DocumentLayout == bbq_code_layout::packed_4b) { + bbq::code_inner_product_packed_4b_symmetric_2x1( + row_a0, row_a1, row_b, total0, total1); + } else if constexpr (DocumentLayout == QueryLayout && + (DocumentLayout == bbq_code_layout::packed_8b || + DocumentLayout == bbq_code_layout::packed_7b)) { + // packed_7b is packed_8b with the top bit masked off, matching code_inner_product's + // (1 << bits) - 1 mask for the same two layouts. + constexpr uint8_t code_mask = DocumentLayout == bbq_code_layout::packed_7b ? 0x7Fu : 0xFFu; + bbq::code_inner_product_packed_8b_2x1( + row_a0, row_a1, row_b, total0, total1, code_mask); + } else { + bbq:: + code_inner_product_planes_2x1( + row_a0, row_a1, row_b, total0, total1); + } +} + +// -------------------------------------------------------------------------- +// Dequantization +// Turns a raw code inner product into a final float distance. +// -------------------------------------------------------------------------- + +// Per-row dequantization terms needed by bbq_calculate_metric. +// row_norm is read directly from the quantizer view instead, +// since register pressure in the BBQ local-join kernels is already tight. +struct bbq_dequant_factors { + float lower; + float delta; + float sum_delta; + float corrections; +}; + +template +__device__ __forceinline__ bbq_dequant_factors +get_dequant_factors(const quantizer_view& quantizer, int64_t row) +{ + return bbq_dequant_factors{quantizer.lower_intervals(row), + quantizer.dequant_delta(row), + quantizer.dequant_sum_delta(row), + quantizer.additional_corrections(row)}; +} + +// Converts one raw BBQ dot product into a final (post-epilogue) float distance, given both +// operands' precomputed dequant factors. Evaluated exactly once per matrix cell +// dim/centroid_norm_sq/row_norm both come directly from quantizer_document/quantizer_query rather +// than being passed separately, since every caller reads them the same id-indexed way -- row_norm +// is only read for CosineExpanded (skipped entirely otherwise). +template +__device__ __forceinline__ float bbq_calculate_metric( + uint32_t raw, + const bbq_dequant_factors& doc_factors, + const bbq_dequant_factors& query_factors, + const quantizer_view& quantizer_document, + const quantizer_view& quantizer_query, + cuvs::distance::DistanceType metric, + DistEpilogue_t dist_epilogue, + Index_t document_id, + Index_t query_id) +{ + constexpr bool can_postprocess_dist = std::is_same_v; + const float dim = static_cast(quantizer_document.dim()); + + const float centered = dim * doc_factors.lower * query_factors.lower + + query_factors.lower * doc_factors.sum_delta + + doc_factors.lower * query_factors.sum_delta + + doc_factors.delta * query_factors.delta * static_cast(raw); + const float corrections = doc_factors.corrections + query_factors.corrections; + float d; + if (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + const float raw_distance = corrections - 2.0f * centered; + d = raw_distance < 0.0f ? 0.0f : raw_distance; + if (!can_postprocess_dist && metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + d = sqrtf(d); + } + } else if (metric == cuvs::distance::DistanceType::InnerProduct) { + d = -(centered + corrections - quantizer_document.centroid_norm_sq); + } else { // CosineExpanded + const float norm_product = + quantizer_document.row_norm(document_id) * quantizer_query.row_norm(query_id); + const float dot = centered + corrections - quantizer_document.centroid_norm_sq; + d = norm_product > 0.0f ? 1.0f - dot / sqrtf(norm_product) : 0.0f; + } + return dist_epilogue(d, document_id, query_id); +} + +#endif // __CUDACC__ + +} // namespace preprocessing::quantize::bbq +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c6a434e4a9..2eed6a3853 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -41,6 +41,7 @@ function(ConfigureTest) ${TEST_NAME} PRIVATE cuvs cuvs::cuvs + cuvs_internal raft::raft rtcx::rtcx GTest::gtest @@ -274,11 +275,19 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME NEIGHBORS_ANN_CAGRA_BBQ_UINT32_TEST + PATH neighbors/ann_cagra/test_bbq_uint32_t.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_NN_DESCENT_TEST PATH neighbors/ann_nn_descent/test_float_uint32_t.cu neighbors/ann_nn_descent/test_int8_t_uint32_t.cu neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu + neighbors/ann_nn_descent/test_bbq_uint32_t.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu new file mode 100644 index 0000000000..f3c0b2a1ff --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_bbq_uint32_t.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_cagra_bbq.cuh" + +#include + +namespace cuvs::neighbors::cagra { + +TEST_P(AnnCagraBbqTest, AnnCagraBbqSearchRecall) { this->testSearchRecall(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqGraphShape) { this->testGraphShape(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqGraphOnlyBuild) { this->testGraphOnlyBuild(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqSerializeRoundTrip) { this->testSerializeRoundTrip(); } +TEST_P(AnnCagraBbqTest, AnnCagraBbqUnsupportedParams) { this->testUnsupportedParams(); } + +INSTANTIATE_TEST_CASE_P(AnnCagraBbqTest, AnnCagraBbqTest, ::testing::ValuesIn(bbq_inputs)); + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_bbq.cuh b/cpp/tests/neighbors/ann_cagra_bbq.cuh new file mode 100644 index 0000000000..9ecc4bd1d4 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra_bbq.cuh @@ -0,0 +1,379 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "ann_cagra.cuh" + +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +struct AnnCagraBbqInputs { + int n_queries; + int n_rows; + int dim; + int k; + int graph_degree; + cuvs::distance::DistanceType metric; + /** Query-side layout; the code width follows from it. */ + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout; + /** + * Document-side layout of a second, coarser quantizer. When set, the dataset carries two + * quantizers and the asymmetric NN-descent path is used. + */ + std::optional second_layout; + /** Fraction of the full-precision build's recall that the BBQ build must retain. */ + double min_recall_ratio; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraBbqInputs& p) +{ + os << "{n_queries=" << p.n_queries << ", n_rows=" << p.n_rows << ", dim=" << p.dim + << ", k=" << p.k << ", graph_degree=" << p.graph_degree << ", metric=" + << cuvs::neighbors::print_metric{static_cast((int)p.metric)} + << ", layout=" << static_cast(p.layout) << ", second_layout=" + << (p.second_layout.has_value() ? static_cast(*p.second_layout) : -1) << "}"; + return os; +} + +class AnnCagraBbqTest : public ::testing::TestWithParam { + public: + AnnCagraBbqTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + search_queries(0, stream_) + { + } + + protected: + /** Quantize the float database on the host and upload the codes. */ + auto quantize_database() -> cuvs::neighbors::device_bbq_dataset + { + std::vector host_data(static_cast(ps.n_rows) * ps.dim); + raft::update_host(host_data.data(), database.data(), host_data.size(), stream_); + raft::resource::sync_stream(handle_); + + return cuvs_internal::bbq::quantize_to_device(handle_, + host_data.data(), + ps.n_rows, + ps.dim, + ps.metric, + ps.layout, + ps.second_layout.value_or(ps.layout)); + } + + [[nodiscard]] auto default_index_params() const -> cagra::index_params + { + cagra::index_params params; + params.metric = ps.metric; + params.graph_degree = ps.graph_degree; + params.intermediate_graph_degree = 2 * ps.graph_degree; + // The BBQ path always uses nn-descent; pin the dense baseline to it as well so the two builds + // differ only in the precision of the distances driving graph construction. + params.graph_build_params = + cagra::graph_build_params::nn_descent_params(params.intermediate_graph_degree, ps.metric); + return params; + } + + /** Brute-force top-k over the current database, as the ground truth for recall. */ + auto naive_neighbours() -> std::vector + { + size_t queries_size = static_cast(ps.n_queries) * ps.k; + std::vector indices_naive(queries_size); + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indices_naive_dev(queries_size, stream_); + cuvs::neighbors::naive_knn(handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + search_queries.data(), + database.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + return indices_naive; + } + + /** Top-k neighbors @p index returns for the test queries. */ + template + auto search_neighbors(IndexT const& index, cagra::search_params search_params = {}) + -> std::vector + { + size_t queries_size = static_cast(ps.n_queries) * ps.k; + rmm::device_uvector distances_dev(queries_size, stream_); + rmm::device_uvector indices_dev(queries_size, stream_); + + auto search_queries_view = raft::make_device_matrix_view( + search_queries.data(), ps.n_queries, ps.dim); + auto indices_out_view = + raft::make_device_matrix_view(indices_dev.data(), ps.n_queries, ps.k); + auto dists_out_view = + raft::make_device_matrix_view(distances_dev.data(), ps.n_queries, ps.k); + + cagra::search( + handle_, search_params, index, search_queries_view, indices_out_view, dists_out_view); + + std::vector indices_cagra(queries_size); + raft::update_host(indices_cagra.data(), indices_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + return indices_cagra; + } + + template + auto search_recall(IndexT const& index, std::vector const& ground_truth) -> double + { + auto [recall, match_count, total_count] = + calc_recall(ground_truth, search_neighbors(index), ps.n_queries, ps.k); + return recall; + } + + /** + * The end-to-end contract: a graph built purely from quantized codes navigates the uncompressed + * vectors nearly as well as one built from the full-precision vectors. Comparing against a dense + * build on the same data and parameters keeps the bar independent of how much recall the search + * configuration itself can reach. + */ + void testSearchRecall() + { + // CAGRA search has no L2SqrtExpanded kernels, so this metric is only exercised by the + // build-side tests below. + if (ps.metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + GTEST_SKIP() << "CAGRA search does not support L2SqrtExpanded"; + } + + auto ground_truth = naive_neighbours(); + auto database_view = + raft::make_device_matrix_view(database.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra device_padded(handle_, + database_view); + + double reference_recall = 0.0; + { + auto dense_index = cagra::build(handle_, default_index_params(), device_padded.view); + dense_index = cagra::update_dataset(handle_, std::move(dense_index), device_padded.view); + reference_recall = search_recall(dense_index, ground_truth); + } + + double bbq_recall = 0.0; + { + // `owning_codes` backs the view held by the built index, so it must outlive it. + auto owning_codes = quantize_database(); + auto graph_index = + cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + ASSERT_EQ(graph_index.graph_size(), static_cast(ps.n_rows)); + + // Rebinding the uncompressed vectors is what makes the BBQ-built graph searchable. + auto index = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); + bbq_recall = search_recall(index, ground_truth); + } + + const double min_recall = ps.min_recall_ratio * reference_recall; + RAFT_LOG_INFO( + "CAGRA BBQ build (layout=%d, second_layout=%d): recall=%f, dense reference=%f, " + "retained=%.1f%% (required >= %.1f%%)", + static_cast(ps.layout), + ps.second_layout.has_value() ? static_cast(*ps.second_layout) : -1, + bbq_recall, + reference_recall, + 100.0 * bbq_recall / reference_recall, + 100.0 * ps.min_recall_ratio); + + // A broken baseline would make the ratio check meaningless. + ASSERT_GT(reference_recall, 0.8) << "the full-precision CAGRA baseline is unexpectedly poor"; + EXPECT_GE(bbq_recall, min_recall) + << "recall " << bbq_recall << " retains only " << 100.0 * bbq_recall / reference_recall + << "% of the full-precision build's " << reference_recall; + } + + /** The optimized graph has the requested shape and refers only to existing rows. */ + void testGraphShape() + { + auto owning_codes = quantize_database(); + auto index = cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + + ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); + // The BBQ dataset view is attached by default, so the index reports the quantized shape. + ASSERT_EQ(index.size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.dim(), static_cast(ps.dim)); + + auto graph_host = raft::make_host_matrix(ps.n_rows, ps.graph_degree); + raft::copy(graph_host.data_handle(), index.graph().data_handle(), graph_host.size(), stream_); + raft::resource::sync_stream(handle_); + + for (int64_t i = 0; i < ps.n_rows; i++) { + for (int64_t j = 0; j < ps.graph_degree; j++) { + ASSERT_LT(graph_host(i, j), static_cast(ps.n_rows)) + << "graph node " << i << " has an out-of-range neighbor at position " << j; + } + } + } + + /** + * A BBQ index file holds the graph alone: the codes and their quantizers live outside the + * index, so the restored graph only becomes searchable once a dataset is reattached, and then + * it answers exactly as the original did. + */ + void testSerializeRoundTrip() + { + if (ps.metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + GTEST_SKIP() << "CAGRA search does not support L2SqrtExpanded"; + } + + auto database_view = + raft::make_device_matrix_view(database.data(), ps.n_rows, ps.dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra device_padded(handle_, + database_view); + + auto owning_codes = quantize_database(); + auto graph_index = + cagra::build(handle_, default_index_params(), owning_codes.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(handle_, stored, graph_index); + + device_bbq_index restored{handle_}; + cagra::deserialize(handle_, stored, &restored); + + ASSERT_EQ(restored.size(), graph_index.size()); + ASSERT_EQ(restored.graph_size(), graph_index.graph_size()); + ASSERT_EQ(restored.graph_degree(), graph_index.graph_degree()); + EXPECT_EQ(restored.metric(), graph_index.metric()); + EXPECT_EQ(restored.dataset().n_rows(), 0); + EXPECT_TRUE(restored.dataset().quantizers.empty()); + + auto original = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); + auto reattached = cagra::update_dataset(handle_, std::move(restored), device_padded.view); + // Use single-CTA kernel to check for exact search results. + cagra::search_params search_params; + search_params.algo = cagra::search_algo::SINGLE_CTA; + EXPECT_EQ(search_neighbors(reattached, search_params), + search_neighbors(original, search_params)); + } + + /** `attach_dataset_on_build = false` yields a graph without any dataset binding. */ + void testGraphOnlyBuild() + { + auto params = default_index_params(); + params.attach_dataset_on_build = false; + + auto owning_codes = quantize_database(); + auto index = cagra::build(handle_, params, owning_codes.as_dataset_view()); + + ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); + ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); + EXPECT_EQ(index.dataset().n_rows(), 0); + EXPECT_TRUE(index.dataset().quantizers.empty()); + } + + /** Only NN-descent graph construction and the four BBQ metrics are accepted. */ + void testUnsupportedParams() + { + auto owning_codes = quantize_database(); + auto dataset = owning_codes.as_dataset_view(); + + auto ivf_pq_params = default_index_params(); + ivf_pq_params.graph_build_params = cagra::graph_build_params::ivf_pq_params( + raft::matrix_extent(ps.n_rows, ps.dim), ps.metric); + EXPECT_THROW(cagra::build(handle_, ivf_pq_params, dataset), raft::exception); + + auto iterative_params = default_index_params(); + iterative_params.graph_build_params = cagra::graph_build_params::iterative_search_params(); + EXPECT_THROW(cagra::build(handle_, iterative_params, dataset), raft::exception); + + auto l1_params = default_index_params(); + l1_params.metric = cuvs::distance::DistanceType::L1; + EXPECT_THROW(cagra::build(handle_, l1_params, dataset), raft::exception); + } + + void SetUp() override + { + // nn-descent rejects packed_4b below sm_75. + if (ps.layout == cuvs::preprocessing::quantize::bbq::bbq_code_layout::packed_4b && + cuvs::neighbors::device_compute_capability() < 75) { + GTEST_SKIP() << "packed_4b requires int4 tensor cores (compute capability 7.5 or newer)"; + } + database.resize(static_cast(ps.n_rows) * ps.dim, stream_); + search_queries.resize(static_cast(ps.n_queries) * ps.dim, stream_); + raft::random::RngState r(1234ULL); + InitDataset(handle_, database.data(), ps.n_rows, ps.dim, ps.metric, r); + InitDataset(handle_, search_queries.data(), ps.n_queries, ps.dim, ps.metric, r); + raft::resource::sync_stream(handle_); + } + + void TearDown() override + { + raft::resource::sync_stream(handle_); + database.resize(0, stream_); + search_queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnCagraBbqInputs ps; + rmm::device_uvector database; + rmm::device_uvector search_queries; +}; + +/** + * Search runs on exact distances, so recall mostly reflects how well a graph built from `bits`-wide + * codes navigates. Quality degrades gracefully as the codes get coarser, so the fraction of the + * full-precision build's recall that must be retained is looser for the narrow codes. + */ +inline const std::vector bbq_inputs = [] { + using cuvs::preprocessing::quantize::bbq::bbq_code_layout; + using opt_layout = std::optional; + const std::vector> code_specs{ + {bbq_code_layout::packed_1b, opt_layout{}, 0.85}, + {bbq_code_layout::transposed_2b, opt_layout{}, 0.92}, + {bbq_code_layout::transposed_2b, opt_layout{bbq_code_layout::packed_1b}, 0.92}, + {bbq_code_layout::packed_4b, opt_layout{}, 0.95}, + {bbq_code_layout::packed_4b, opt_layout{bbq_code_layout::packed_1b}, 0.90}, + {bbq_code_layout::transposed_4b, opt_layout{bbq_code_layout::packed_1b}, 0.90}, + {bbq_code_layout::transposed_4b, opt_layout{bbq_code_layout::transposed_2b}, 0.88}, + {bbq_code_layout::packed_7b, opt_layout{}, 0.95}, + {bbq_code_layout::packed_8b, opt_layout{}, 0.95}}; + + std::vector out; + for (const auto& [layout, second_layout, min_recall_ratio] : code_specs) { + const auto batch = raft::util::itertools::product( + {200}, // n_queries + {4000}, // n_rows + {128}, // dim + {10}, // k + {32}, // graph_degree + {cuvs::distance::DistanceType::L2Expanded, + // cuvs::distance::DistanceType::L2SqrtExpanded, + cuvs::distance::DistanceType::InnerProduct, + cuvs::distance::DistanceType::CosineExpanded}, + {layout}, + {second_layout}, + {min_recall_ratio}); + out.insert(out.end(), batch.begin(), batch.end()); + } + return out; +}(); + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu b/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu new file mode 100644 index 0000000000..bea903ef49 --- /dev/null +++ b/cpp/tests/neighbors/ann_nn_descent/test_bbq_uint32_t.cu @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_nn_descent_bbq.cuh" + +#include + +namespace cuvs::neighbors::nn_descent { +TEST_P(AnnNNDescentBbqTest, AnnNNDescentBbq) { this->testNNDescent(); } + +INSTANTIATE_TEST_CASE_P(AnnNNDescentBbqTest, AnnNNDescentBbqTest, ::testing::ValuesIn(bbq_inputs)); +} // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_nn_descent_bbq.cuh b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh new file mode 100644 index 0000000000..295cae935f --- /dev/null +++ b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "ann_nn_descent.cuh" + +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::nn_descent { +// The host reference quantizer is shared with the ann-bench CAGRA wrapper, so these tests and +// the benchmark can never disagree about the code format. +namespace cpu_bbq = cuvs_internal::bbq; +using cuvs::preprocessing::quantize::bbq::get_bit_width; +using cuvs_internal::bbq::make_device_bbq_dataset; + +struct AnnNNDescentBbqInputs : AnnNNDescentInputs { + cuvs::preprocessing::quantize::bbq::bbq_code_layout layout; + std::optional second_dataset_layout; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnNNDescentBbqInputs& p) +{ + os << "dataset shape=" << p.n_rows << "x" << p.dim << ", graph_degree=" << p.graph_degree + << ", metric=" + << cuvs::neighbors::print_metric{static_cast((int)p.metric)} + << (p.host_dataset ? ", host" : ", device") << ", layout=" << static_cast(p.layout) + << ", second_dataset_layout=" + << (p.second_dataset_layout.has_value() ? static_cast(p.second_dataset_layout.value()) + : -1) + << std::endl; + return os; +} + +class AnnNNDescentBbqTest : public ::testing::TestWithParam { + public: + AnnNNDescentBbqTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(raft::make_device_matrix(handle_, ps.n_rows, ps.dim)) + { + } + + protected: + void testNNDescent() + { + // nn-descent rejects packed_4b below sm_75. + if (ps.layout == cuvs::preprocessing::quantize::bbq::bbq_code_layout::packed_4b && + cuvs::neighbors::device_compute_capability() < 75) { + GTEST_SKIP() << "packed_4b requires int4 tensor cores (compute capability 7.5 or newer)"; + } + if (ps.second_dataset_layout.has_value()) { + // The document must be strictly coarser than the query; the pair's layouts are stated in + // the spec, so validity of the layout combination is the spec's business, not inferred here. + const uint32_t query_bits = get_bit_width(ps.layout); + const uint32_t second_bits = get_bit_width(ps.second_dataset_layout.value()); + if (query_bits > 4 || query_bits == second_bits || query_bits == 1) { + GTEST_SKIP() << "Second dataset is N/A: bits=" << query_bits + << ", layout=" << static_cast(ps.layout) + << " and second bits=" << second_bits; + } + } + size_t queries_size = ps.n_rows * ps.graph_degree; + std::vector indices_NNDescent(queries_size); + std::vector distances_NNDescent(queries_size); + std::vector indices_naive(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indices_naive_dev(queries_size, stream_); + naive_knn(handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + database.data_handle(), + database.data_handle(), + ps.n_rows, + ps.n_rows, + ps.dim, + ps.graph_degree, + ps.metric); + raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + { + std::vector host_data(static_cast(ps.n_rows) * ps.dim); + raft::update_host(host_data.data(), database.data_handle(), host_data.size(), stream_); + raft::resource::sync_stream(handle_); + + auto owning_dataset = + cuvs_internal::bbq::quantize_to_device(handle_, + host_data.data(), + ps.n_rows, + ps.dim, + ps.metric, + ps.layout, + ps.second_dataset_layout.value_or(ps.layout)); + auto dataset = owning_dataset.as_dataset_view(); + nn_descent::index_params index_params; + index_params.metric = ps.metric; + index_params.graph_degree = ps.graph_degree; + index_params.intermediate_graph_degree = 2 * ps.graph_degree; + index_params.max_iterations = 100; + index_params.return_distances = true; + + auto index = nn_descent::build(handle_, index_params, dataset); + + raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_); + ASSERT_TRUE(index.distances().has_value()); + raft::copy( + distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + EXPECT_TRUE(eval_neighbours(indices_naive, + indices_NNDescent, + distances_naive, + distances_NNDescent, + ps.n_rows, + ps.graph_degree, + 0.001, + ps.min_recall)); + } + + void SetUp() override + { + raft::random::RngState r(1234ULL); + raft::random::normal(handle_, r, database.data_handle(), ps.n_rows * ps.dim, 0.1f, 2.0f); + raft::resource::sync_stream(handle_); + } + + void TearDown() override { raft::resource::sync_stream(handle_); } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnNNDescentBbqInputs ps; + raft::device_matrix database; +}; + +// Estimated recall based on bruteforce (InnerProduct): 1: 0.23, 2: 0.52, 4: 0.85, 7: 0.98, 8: 0.99. +const std::vector bbq_inputs = [] { + using cuvs::preprocessing::quantize::bbq::bbq_code_layout; + const std::vector>> + code_specifications{ + // min_recall, query layout, document layout (both widths follow from the layouts) + {0.15, bbq_code_layout::packed_1b, std::optional{}}, + {0.50, bbq_code_layout::transposed_2b, std::optional{}}, + {0.27, + bbq_code_layout::transposed_2b, + std::optional{bbq_code_layout::packed_1b}}, + {0.80, bbq_code_layout::packed_4b, std::optional{}}, + // Asymmetric packed_4b queries take the int4 wmma path (SelfJoin = false). packed_1b is the + // only document layout that promotes to it -- transposed_2b never reaches this kernel. At + // dim=256 this is also the only coverage of the phase-2 staging skip (n_tiles == 1) outside + // a self-join. + {0.35, + bbq_code_layout::packed_4b, + std::optional{bbq_code_layout::packed_1b}}, + // Asymmetric transposed_4b queries (1 + 4t, 2t + 4t) take the SIMT path. + {0.35, + bbq_code_layout::transposed_4b, + std::optional{bbq_code_layout::packed_1b}}, + {0.65, + bbq_code_layout::transposed_4b, + std::optional{bbq_code_layout::transposed_2b}}, + {0.80, bbq_code_layout::packed_7b, std::optional{}}, + {0.80, bbq_code_layout::packed_8b, std::optional{}}}; + std::vector out; + for (const auto& [min_recall, layout, second_layout] : code_specifications) { + const auto batch = raft::util::itertools::product( + {2000}, + {256}, // dim + {64}, // graph_degree + {cuvs::distance::DistanceType::L2Expanded, + cuvs::distance::DistanceType::L2SqrtExpanded, + cuvs::distance::DistanceType::InnerProduct, + cuvs::distance::DistanceType::CosineExpanded}, + {false}, // host_dataset + {min_recall}, + {layout}, + {second_layout}); + out.insert(out.end(), batch.begin(), batch.end()); + } + return out; +}(); + +} // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_utils.cuh b/cpp/tests/neighbors/ann_utils.cuh index a45628ed62..c1022eba17 100644 --- a/cpp/tests/neighbors/ann_utils.cuh +++ b/cpp/tests/neighbors/ann_utils.cuh @@ -28,6 +28,13 @@ namespace cuvs::neighbors { +/** Compute capability of the current device as major * 10 + minor. */ +inline auto device_compute_capability() -> int +{ + auto [major, minor] = raft::getComputeCapability(); + return major * 10 + minor; +} + struct print_dtype { cudaDataType_t value; }; diff --git a/fern/docs.yml b/fern/docs.yml index ca3ea6e7df..1daa52cf02 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -314,6 +314,8 @@ navigation: path: "./pages/c_api/c-api-neighbors-tiered-index.md" - page: "Neighbors Vamana" path: "./pages/c_api/c-api-neighbors-vamana.md" + - page: "Preprocessing Quantize Bbq" + path: "./pages/c_api/c-api-preprocessing-quantize-bbq.md" - page: "Preprocessing Quantize Binary" path: "./pages/c_api/c-api-preprocessing-quantize-binary.md" - page: "Preprocessing PCA" @@ -398,6 +400,8 @@ navigation: path: "./pages/cpp_api/cpp-api-neighbors-tiered-index.md" - page: "Neighbors Vamana" path: "./pages/cpp_api/cpp-api-neighbors-vamana.md" + - page: "Preprocessing Quantize Bbq" + path: "./pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md" - page: "Preprocessing Quantize Binary" path: "./pages/cpp_api/cpp-api-preprocessing-quantize-binary.md" - page: "Preprocessing PCA" @@ -472,6 +476,8 @@ navigation: - section: "Java API Documentation" path: "./pages/java_api/index.md" contents: + - page: "BbqQuantizer" + path: "./pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md" - page: "CuVSDeviceMatrix" path: "./pages/java_api/java-api-com-nvidia-cuvs-cuvsdevicematrix.md" - page: "CuVSHostMatrix" diff --git a/fern/pages/c_api/c-api-core-dataset.md b/fern/pages/c_api/c-api-core-dataset.md index 915c16cbbc..9d989183d6 100644 --- a/fern/pages/c_api/c-api-core-dataset.md +++ b/fern/pages/c_api/c-api-core-dataset.md @@ -17,7 +17,8 @@ Generic dataset layout kind for C API dataset handles. typedef enum { CUVS_DATASET_LAYOUT_STANDARD = 0, CUVS_DATASET_LAYOUT_PADDED = 1, - CUVS_DATASET_LAYOUT_PQ = 2 + CUVS_DATASET_LAYOUT_PQ = 2, + CUVS_DATASET_LAYOUT_BBQ = 3 } cuvsDatasetLayout_t; ``` @@ -28,6 +29,7 @@ typedef enum { | `CUVS_DATASET_LAYOUT_STANDARD` | `0` | | `CUVS_DATASET_LAYOUT_PADDED` | `1` | | `CUVS_DATASET_LAYOUT_PQ` | `2` | +| `CUVS_DATASET_LAYOUT_BBQ` | `3` | ### cuvsDatasetMemType_t diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md index 7ac45afa5a..13c663a7c8 100644 --- a/fern/pages/c_api/c-api-neighbors-cagra.md +++ b/fern/pages/c_api/c-api-neighbors-cagra.md @@ -827,11 +827,11 @@ cuvsDataset_t dataset, cuvsCagraIndex_t index); ``` -The memory space and layout `dataset` was constructed with select the C++ build overload. Build the handle with an owning factory or the matching dataset view factory (`cuvsDatasetMakePaddedView` / `cuvsDatasetMakeStandardView`). +The memory space and layout `dataset` was constructed with select the C++ build overload. Build the handle with an owning factory or the matching dataset view factory (`cuvsDatasetMakePaddedView`, `cuvsDatasetMakeStandardView`, or `cuvsDatasetMakeBbqView`). Note that a dataset residing in host memory produces a host-backed index, which must be made search-ready with `cuvsCagraUpdateDataset` (using a device-padded dataset) before calling `cuvsCagraSearch`. -A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the VPQ dataset must outlive the index because the index stores a non-owning view. +A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iterative CAGRA-Q index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the VPQ dataset must outlive the index because the index stores a non-owning view. A `CUVS_DATASET_LAYOUT_BBQ` dataset builds a graph-only index; attach a searchable dataset with `cuvsCagraUpdateDataset` before search. **Parameters** @@ -935,7 +935,7 @@ const char* filename, cuvsCagraIndex_t index); ``` -This supports dense and PQ-backed indexes. The dataset must be attached separately after loading the graph. +This supports dense, PQ-backed, and BBQ-built indexes. The serialized file does not contain vector data. After deserialization the index cannot be searched until a compatible dataset is attached with `cuvsCagraUpdateDataset`. Experimental, both the API and the serialization format are subject to change. @@ -962,7 +962,7 @@ const char* filename, cuvsCagraIndex_t index); ``` -The index stores a non-owning dataset view. The caller must keep the dataset backing that view alive while this function runs. Returns CUVS_ERROR without modifying the destination file if the index has no attached dataset. PQ datasets are not serialized by this function. +The index stores a non-owning dataset view. The caller must keep the memory of the dataset backing that view alive while this function runs. Returns CUVS_ERROR without modifying the destination file if the index has no attached dataset. PQ and BBQ datasets are not serialized by this function. Experimental, both the API and the serialization format are subject to change. diff --git a/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md b/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md new file mode 100644 index 0000000000..8f57b980c7 --- /dev/null +++ b/fern/pages/c_api/c-api-preprocessing-quantize-bbq.md @@ -0,0 +1,153 @@ +--- +slug: api-reference/c-api-preprocessing-quantize-bbq +--- + +# Bbq + +_Source header: `cuvs/preprocessing/quantize/bbq.h`_ + +## C API for Better Binary Quantization datasets + + +### cuvsBbqCodeLayout_t + +Storage layout of BBQ/OSQ quantized component codes in each dataset row. + +```c +typedef enum { + CUVS_BBQ_CODE_LAYOUT_PACKED_1B = 0, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B, + CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_4B, + CUVS_BBQ_CODE_LAYOUT_PACKED_7B, + CUVS_BBQ_CODE_LAYOUT_PACKED_8B +} cuvsBbqCodeLayout_t; +``` + +**Values** + +| Name | Value | Description | +| --- | --- | --- | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_1B` | `0` | Each dimension is quantized to a single bit and packed into bytes. Reflects Lucene's OptimizedScalarQuantizer.packAsBinary. | +| `CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_2B` | `` | Each dimension is quantized to 2 bits, stored as 2 bitplanes. Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only (paired with a transposed_4b or packed_1b operand); | +| `CUVS_BBQ_CODE_LAYOUT_TRANSPOSED_4B` | `` | Each dimension is quantized to 4 bits, optimized for bitwise operations. Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of every dimension is in the first set dimensions bits, or (dimensions/8) bytes. The second, third, and fourth bits are in the second, third, and fourth set of dimensions bits, respectively. Format used for queries. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_4B` | `` | Each dimension is quantized to 4 bits, two values are packed into each output byte. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_7B` | `` | Each dimension is quantized to 7 bits and treated as a signed value. | +| `CUVS_BBQ_CODE_LAYOUT_PACKED_8B` | `` | Each dimension is quantized to 8 bits and treated as an unsigned value. | + + +### cuvsBbqQuantizer + +Better Binary Quantization ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch exposes it as BBQ. + +BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared centroid, no random rotation, and OSQ codes. + +RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain additional per-vector information so search is better than naïve sign-bit comparison. They differ in transformation and scale representation. RaBitQ commonly separates residual magnitude from direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector scalar intervals to interpret the compressed residual codes. + +```c +typedef struct cuvsBbqQuantizer { + uintptr_t addr; + DLDataType dtype; + bool is_owning; +} cuvsBbqQuantizer; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `addr` | `uintptr_t` | | +| `dtype` | `DLDataType` | | +| `is_owning` | `bool` | | + + +### cuvsBbqQuantizerCreateView + +Create a BBQ quantizer view from caller-owned device tensors. + +```c +cuvsError_t cuvsBbqQuantizerCreateView( +DLManagedTensor* codes, +DLManagedTensor* lower_intervals, +DLManagedTensor* upper_intervals, +DLManagedTensor* additional_corrections, +DLManagedTensor* quantized_component_sums, +DLManagedTensor* centroid, +DLManagedTensor* dequant_delta, +DLManagedTensor* dequant_sum_delta, +DLManagedTensor* row_norm, +cuvsBbqCodeLayout_t layout, +cuvsDistanceType metric, +float centroid_norm_sq, +cuvsBbqQuantizer_t* quantizer); +``` + +Tensors are not copied and must remain valid while a derived dataset is in use. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `codes` | in | `DLManagedTensor*` | uint8 device matrix containing encoded rows | +| `lower_intervals` | in | `DLManagedTensor*` | float32 device vector with one lower interval per row | +| `upper_intervals` | in | `DLManagedTensor*` | float32 device vector with one upper interval per row | +| `additional_corrections` | in | `DLManagedTensor*` | float32 device vector with one correction per row | +| `quantized_component_sums` | in | `DLManagedTensor*` | int32 device vector with one component sum per row | +| `centroid` | in | `DLManagedTensor*` | device vector containing the dataset centroid | +| `dequant_delta` | in | `DLManagedTensor*` | float32 device vector with one dequantization delta per row | +| `dequant_sum_delta` | in | `DLManagedTensor*` | float32 device vector with one delta-times-sum value per row | +| `row_norm` | in | `DLManagedTensor*` | float32 device vector with one original-space squared norm per row | +| `layout` | in | [`cuvsBbqCodeLayout_t`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqcodelayout-t) | encoded code layout | +| `metric` | in | [`cuvsDistanceType`](/api-reference/c-api-distance-distance#cuvsdistancetype) | distance metric associated with the encoded dataset | +| `centroid_norm_sq` | in | `float` | squared norm of the centroid | +| `quantizer` | out | [`cuvsBbqQuantizer_t*`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | newly allocated non-owning quantizer handle | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + + +### cuvsBbqQuantizerDestroy + +Destroy a BBQ quantizer without destroying its caller-owned tensors. + +```c +cuvsError_t cuvsBbqQuantizerDestroy(cuvsBbqQuantizer_t quantizer); +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `quantizer` | in | [`cuvsBbqQuantizer_t`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | quantizer handle to destroy | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + + +### cuvsDatasetMakeBbqView + +Create a non-owning device BBQ dataset view. + +```c +cuvsError_t cuvsDatasetMakeBbqView(cuvsResources_t res, +cuvsBbqQuantizer_t* quantizers, +size_t num_quantizers, +cuvsDataset_t* dataset); +``` + +Accepts one symmetric quantizer or two compatible asymmetric quantizers. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | [`cuvsResources_t`](/api-reference/c-api-core-c-api#cuvsresources-t) | cuVS resources | +| `quantizers` | in | [`cuvsBbqQuantizer_t*`](/api-reference/c-api-preprocessing-quantize-bbq#cuvsbbqquantizer) | array containing one or two BBQ quantizer handles | +| `num_quantizers` | in | `size_t` | number of elements in `quantizers` | +| `dataset` | out | `cuvsDataset_t*` | newly allocated non-owning BBQ dataset handle | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) diff --git a/fern/pages/c_api/index.md b/fern/pages/c_api/index.md index 968ef01b0e..cd70102764 100644 --- a/fern/pages/c_api/index.md +++ b/fern/pages/c_api/index.md @@ -41,6 +41,7 @@ These pages are generated from the documented public headers in the cuVS source ## Preprocessing +- [Bbq](/api-reference/c-api-preprocessing-quantize-bbq) - [Binary](/api-reference/c-api-preprocessing-quantize-binary) - [PCA](/api-reference/c-api-preprocessing-pca) - [PQ](/api-reference/c-api-preprocessing-quantize-pq) diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 9097f4cb47..2560f59ef1 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -1098,6 +1098,104 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) built `host_standard_index<uint8_t, uint32_t>` +**Additional overload:** `neighbors::cagra::build` + +Build from a device BBQ-quantized dataset view. + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +The kNN graph is built from the quantized codes alone, so the uncompressed vectors are never needed and peak memory is driven by the code size. Only nn-descent graph construction is available (IVF-PQ, iterative CAGRA search, and ACE all read uncompressed vectors), and the metric must be one of L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct and must match the metric the quantizer corrections were generated for. + +The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the type-changing `update_dataset` with an uncompressed device-padded dataset to search the resulting graph. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | raft resources | +| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view const&` | device BBQ dataset view [n_rows, dim] | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +built `device_bbq_index<float, uint32_t>` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + +**Additional overload:** `neighbors::cagra::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) + +```cpp +auto build(raft::resources const& res, +const cuvs::neighbors::cagra::index_params& params, +cuvs::neighbors::device_bbq_dataset_view const& dataset) +-> cuvs::neighbors::cagra::device_bbq_index; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | `const cuvs::neighbors::cagra::index_params&` | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | + +**Returns** + +`cuvs::neighbors::cagra::device_bbq_index` + ## CAGRA extend functions diff --git a/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md b/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md index abbb9cbe58..396966c202 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-nn-descent.md @@ -528,3 +528,111 @@ Usage example: **Returns** [`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +Build an NN-Descent index directly from a device-resident BBQ dataset. + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +Symmetric compressed-code distances are used during graph construction. Supported metrics are L2Expanded, L2SqrtExpanded, CosineExpanded, and InnerProduct. The dataset's correction terms must have been generated for the selected metric. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | raft resources | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | NN-Descent build parameters | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | BBQ codes and correction terms in device memory | +| `graph` | | `std::optional>` | optional caller-owned host graph
Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) + +**Additional overload:** `neighbors::nn_descent::build` + +cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset, + +```cpp +auto build(raft::resources const& res, +index_params const& params, +cuvs::neighbors::device_bbq_dataset_view dataset, +std::optional> graph = +std::nullopt) -> cuvs::neighbors::nn_descent::index; +``` + +std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t> + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `params` | | [`index_params const&`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index-params) | | +| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `graph` | | `std::optional>` | Default: `std::nullopt`. | + +**Returns** + +[`cuvs::neighbors::nn_descent::index`](/api-reference/cpp-api-neighbors-nn-descent#neighbors-nn-descent-index) diff --git a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md new file mode 100644 index 0000000000..67fa8291c1 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md @@ -0,0 +1,153 @@ +--- +slug: api-reference/cpp-api-preprocessing-quantize-bbq +--- + +# Bbq + +_Source header: `cuvs/preprocessing/quantize/bbq.hpp`_ + +## Better Binary Quantization utilities + + +### preprocessing::quantize::bbq::bbq_code_layout + +Storage layout of BBQ/OSQ quantized component codes in each dataset row. + +```cpp +enum class bbq_code_layout { + packed_1b, + transposed_2b, + transposed_4b, + packed_4b, + packed_7b, + packed_8b +}; +``` + +**Values** + +| Name | Value | Description | +| --- | --- | --- | +| `packed_1b` | `` | Each dimension is quantized to a single bit and packed into bytes. Reflects Lucene's OptimizedScalarQuantizer.packAsBinary. | +| `transposed_2b` | `` | Each dimension is quantized to 2 bits, stored as 2 bitplanes. Reflects Lucene's OptimizedScalarQuantizer.transposeDibit. SIMT popc path only (paired with a transposed_4b or packed_1b operand); | +| `transposed_4b` | `` | Each dimension is quantized to 4 bits, optimized for bitwise operations. Reflects Lucene's OptimizedScalarQuantizer.transposeHalfByte. the first bit of every dimension is in the first set dimensions bits, or (dimensions/8) bytes. The second, third, and fourth bits are in the second, third, and fourth set of dimensions bits, respectively. Format used for queries. | +| `packed_4b` | `` | Each dimension is quantized to 4 bits, two values are packed into each output byte. | +| `packed_7b` | `` | Each dimension is quantized to 7 bits and treated as a signed value. | +| `packed_8b` | `` | Each dimension is quantized to 8 bits and treated as an unsigned value. | + + +### preprocessing::quantize::bbq::get_bit_width + +Bit width of a layout. + +```cpp +constexpr auto get_bit_width(bbq_code_layout layout) noexcept -> uint32_t; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`uint32_t` + + +### preprocessing::quantize::bbq::get_encoded_row_length + +Bytes one row of `dim` components occupies once encoded in `layout`. + +```cpp +constexpr auto get_encoded_row_length(uint32_t dim, bbq_code_layout layout) noexcept -> uint32_t; +``` + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `dim` | | `uint32_t` | | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`uint32_t` + + +### preprocessing::quantize::bbq::quantizer + +Better Binary Quantization ([BBQ](https://www.elastic.co/search-labs/blog/better-binary-quantization-lucene-elasticsearch)) is a vector-quantization approach used in Elasticsearch and Apache Lucene. It builds on ideas introduced in RaBitQ([Gao and Long](https://arxiv.org/pdf/2405.12497, [Gao et al.](https://arxiv.org/pdf/2409.09913)): residual binary codes around a centroid, corrective factors, and efficient bitwise comparison of codes at different bit widths. Lucene implements this as optimized scalar quantization (OSQ) with packed and bit-plane layouts; Elasticsearch exposes it as BBQ. + +BBQ in cuVS designed to be compatible with the Lucene/Elasticsearch dataset: a single shared centroid, no random rotation, and OSQ codes. + +RaBitQ and BBQ in cuVS both compress centroid-relative vectors to low-bit codes and retain additional per-vector information so search is better than naïve sign-bit comparison. They differ in transformation and scale representation. RaBitQ commonly separates residual magnitude from direction, then applies a random orthogonal rotation before binary coding; BBQ uses per-vector scalar intervals to interpret the compressed residual codes. + +```cpp +template +struct quantizer { + raft::device_vector dequant_delta; + raft::device_vector dequant_sum_delta; + raft::device_vector row_norm; + raft::device_matrix codes; + raft::device_vector lower_intervals; + raft::device_vector upper_intervals; + raft::device_vector additional_corrections; + raft::device_vector quantized_component_sums; + raft::device_vector centroid; + bbq_code_layout layout; + cuvs::distance::DistanceType metric; + float centroid_norm_sq; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dequant_delta` | `raft::device_vector` | Precomputed per-row dequantization factors, derived once (offline) from lower/upper_intervals and quantized_component_sums: dequant_delta = (upper-lower)/(2^bits-1) | +| `dequant_sum_delta` | `raft::device_vector` | Precomputed per-row dequantization factors, derived once (offline) from dequant_delta and quantized_component_sums: dequant_sum_delta = dequant_delta * quantized_component_sums. | +| `row_norm` | `raft::device_vector` | Squared norm of the row in original (un-centered) vector space, \|\|x\|\|^2 | +| `codes` | `raft::device_matrix` | | +| `lower_intervals` | `raft::device_vector` | | +| `upper_intervals` | `raft::device_vector` | | +| `additional_corrections` | `raft::device_vector` | | +| `quantized_component_sums` | `raft::device_vector` | | +| `centroid` | `raft::device_vector` | | +| `layout` | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | +| `metric` | [`cuvs::distance::DistanceType`](/api-reference/cpp-api-distance-distance#distance-distancetype) | | +| `centroid_norm_sq` | `float` | | + + +### preprocessing::quantize::bbq::helpers::resolve_dequant_factors + +Derives dequant_delta from lower/upper_intervals and the layout's code width, and + +```cpp +void resolve_dequant_factors( +raft::resources const& res, +raft::device_vector_view dequant_delta, +raft::device_vector_view dequant_sum_delta, +raft::device_vector_view lower_intervals, +raft::device_vector_view upper_intervals, +raft::device_vector_view quantized_component_sums, +bbq_code_layout layout); +``` + +dequant_sum_delta from that delta and quantized_component_sums. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | | `raft::resources const&` | | +| `dequant_delta` | | `raft::device_vector_view` | | +| `dequant_sum_delta` | | `raft::device_vector_view` | | +| `lower_intervals` | | `raft::device_vector_view` | | +| `upper_intervals` | | `raft::device_vector_view` | | +| `quantized_component_sums` | | `raft::device_vector_view` | | +| `layout` | | [`bbq_code_layout`](/api-reference/cpp-api-preprocessing-quantize-bbq#preprocessing-quantize-bbq-bbq-code-layout) | | + +**Returns** + +`void` diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 3c489df78b..07115a5eca 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -43,6 +43,7 @@ These pages are generated from the documented public headers in the cuVS source ## Preprocessing +- [Bbq](/api-reference/cpp-api-preprocessing-quantize-bbq) - [Binary](/api-reference/cpp-api-preprocessing-quantize-binary) - [PCA](/api-reference/cpp-api-preprocessing-pca) - [PQ](/api-reference/cpp-api-preprocessing-quantize-pq) diff --git a/fern/pages/java_api/index.md b/fern/pages/java_api/index.md index 42a495c23f..f94a35fb89 100644 --- a/fern/pages/java_api/index.md +++ b/fern/pages/java_api/index.md @@ -6,6 +6,7 @@ For the Apache Lucene codecs built on this API, see the [Lucene API Documentatio ## Common +- [BbqQuantizer](/api-reference/java-api-com-nvidia-cuvs-bbqquantizer) - [CuVSDeviceMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvsdevicematrix) - [CuVSHostMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvshostmatrix) - [CuVSMatrix](/api-reference/java-api-com-nvidia-cuvs-cuvsmatrix) diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md new file mode 100644 index 0000000000..5f47befedf --- /dev/null +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-bbqquantizer.md @@ -0,0 +1,31 @@ +--- +slug: api-reference/java-api-com-nvidia-cuvs-bbqquantizer +--- + +# BbqQuantizer + +_Java package: `com.nvidia.cuvs`_ + +```java +public final class BbqQuantizer +``` + +Caller-owned tensors describing one encoded BBQ dataset representation. + +The index stores views over these matrices rather than copying them, so they must stay open +for as long as any index built from them is in use. + +## Public Members + +### withCentroidNormSq + +```java +public Builder withCentroidNormSq(float value) +``` + +Sets the squared L2 norm of the centroid. Required: there is no meaningful default, and +leaving it at zero skews inner-product and cosine distances without reporting an error. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java:185`_ + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java:15`_ diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md index 0484513500..b26f36e2d1 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md @@ -613,6 +613,24 @@ an instance of this Builder _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:433`_ +### withBbqDataset + +```java +Builder withBbqDataset(BbqQuantizer... quantizers) +``` + +Builds the graph from one or two encoded BBQ representations. An optional dense dataset +supplied with `#withDataset(CuVSMatrix)` is attached before search; otherwise call +`CagraIndex#updateDataset(PaddedDatasetView)` or +`CagraIndex#updateDataset(PaddedDataset)` before searching. + +The index stores views over the quantizer tensors rather than copying them, so they must +stay open for as long as the index is in use. A dense dataset passed to +`#withDataset(CuVSMatrix)` is owned by the index, as it is for a non-BBQ build, and is +closed with it. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:446`_ + ### withIndexParams ```java @@ -632,7 +650,7 @@ Builder. An instance of this Builder. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:442`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:455`_ ### build @@ -646,6 +664,6 @@ Builds and returns an instance of CagraIndex. an instance of CagraIndex -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:449`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:462`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:26`_ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java new file mode 100644 index 0000000000..04c193cd47 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BbqQuantizer.java @@ -0,0 +1,195 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import java.util.Objects; + +/** + * Caller-owned tensors describing one encoded BBQ dataset representation. + * + *

The index stores views over these matrices rather than copying them, so they must stay open + * for as long as any index built from them is in use. + */ +public final class BbqQuantizer { + public enum CodeLayout { + PACKED_1B(0), + TRANSPOSED_2B(1), + TRANSPOSED_4B(2), + PACKED_4B(3), + PACKED_7B(4), + PACKED_8B(5); + + public final int value; + + CodeLayout(int value) { + this.value = value; + } + } + + private final CuVSMatrix codes; + private final CuVSMatrix lowerIntervals; + private final CuVSMatrix upperIntervals; + private final CuVSMatrix additionalCorrections; + private final CuVSMatrix quantizedComponentSums; + private final CuVSMatrix centroid; + private final CuVSMatrix dequantDelta; + private final CuVSMatrix dequantSumDelta; + private final CuVSMatrix rowNorm; + private final CodeLayout layout; + private final CagraIndexParams.CuvsDistanceType metric; + private final float centroidNormSq; + + private BbqQuantizer(Builder builder) { + codes = Objects.requireNonNull(builder.codes, "codes"); + lowerIntervals = Objects.requireNonNull(builder.lowerIntervals, "lowerIntervals"); + upperIntervals = Objects.requireNonNull(builder.upperIntervals, "upperIntervals"); + additionalCorrections = + Objects.requireNonNull(builder.additionalCorrections, "additionalCorrections"); + quantizedComponentSums = + Objects.requireNonNull(builder.quantizedComponentSums, "quantizedComponentSums"); + centroid = Objects.requireNonNull(builder.centroid, "centroid"); + dequantDelta = Objects.requireNonNull(builder.dequantDelta, "dequantDelta"); + dequantSumDelta = Objects.requireNonNull(builder.dequantSumDelta, "dequantSumDelta"); + rowNorm = Objects.requireNonNull(builder.rowNorm, "rowNorm"); + layout = Objects.requireNonNull(builder.layout, "layout"); + metric = Objects.requireNonNull(builder.metric, "metric"); + centroidNormSq = Objects.requireNonNull(builder.centroidNormSq, "centroidNormSq"); + if (!Float.isFinite(centroidNormSq) || centroidNormSq < 0.0f) { + throw new IllegalArgumentException( + "centroidNormSq must be a finite, non-negative squared norm, but was " + centroidNormSq); + } + } + + public CuVSMatrix getCodes() { + return codes; + } + + public CuVSMatrix getLowerIntervals() { + return lowerIntervals; + } + + public CuVSMatrix getUpperIntervals() { + return upperIntervals; + } + + public CuVSMatrix getAdditionalCorrections() { + return additionalCorrections; + } + + public CuVSMatrix getQuantizedComponentSums() { + return quantizedComponentSums; + } + + public CuVSMatrix getCentroid() { + return centroid; + } + + public CuVSMatrix getDequantDelta() { + return dequantDelta; + } + + public CuVSMatrix getDequantSumDelta() { + return dequantSumDelta; + } + + public CuVSMatrix getRowNorm() { + return rowNorm; + } + + public CodeLayout getLayout() { + return layout; + } + + public CagraIndexParams.CuvsDistanceType getMetric() { + return metric; + } + + public float getCentroidNormSq() { + return centroidNormSq; + } + + public static final class Builder { + private CuVSMatrix codes; + private CuVSMatrix lowerIntervals; + private CuVSMatrix upperIntervals; + private CuVSMatrix additionalCorrections; + private CuVSMatrix quantizedComponentSums; + private CuVSMatrix centroid; + private CuVSMatrix dequantDelta; + private CuVSMatrix dequantSumDelta; + private CuVSMatrix rowNorm; + private CodeLayout layout; + private CagraIndexParams.CuvsDistanceType metric; + private Float centroidNormSq; + + public Builder withCodes(CuVSMatrix value) { + codes = value; + return this; + } + + public Builder withLowerIntervals(CuVSMatrix value) { + lowerIntervals = value; + return this; + } + + public Builder withUpperIntervals(CuVSMatrix value) { + upperIntervals = value; + return this; + } + + public Builder withAdditionalCorrections(CuVSMatrix value) { + additionalCorrections = value; + return this; + } + + public Builder withQuantizedComponentSums(CuVSMatrix value) { + quantizedComponentSums = value; + return this; + } + + public Builder withCentroid(CuVSMatrix value) { + centroid = value; + return this; + } + + public Builder withDequantDelta(CuVSMatrix value) { + dequantDelta = value; + return this; + } + + public Builder withDequantSumDelta(CuVSMatrix value) { + dequantSumDelta = value; + return this; + } + + public Builder withRowNorm(CuVSMatrix value) { + rowNorm = value; + return this; + } + + public Builder withLayout(CodeLayout value) { + layout = value; + return this; + } + + public Builder withMetric(CagraIndexParams.CuvsDistanceType value) { + metric = value; + return this; + } + + /** + * Sets the squared L2 norm of the centroid. Required: there is no meaningful default, and + * leaving it at zero skews inner-product and cosine distances without reporting an error. + */ + public Builder withCentroidNormSq(float value) { + centroidNormSq = value; + return this; + } + + public BbqQuantizer build() { + return new BbqQuantizer(this); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index d9472c001d..5320afe4af 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -433,6 +433,19 @@ interface Builder { */ Builder withDataset(CuVSMatrix dataset); + /** + * Builds the graph from one or two encoded BBQ representations. An optional dense dataset + * supplied with {@link #withDataset(CuVSMatrix)} is attached before search; otherwise call + * {@link CagraIndex#updateDataset(PaddedDatasetView)} or + * {@link CagraIndex#updateDataset(PaddedDataset)} before searching. + * + *

The index stores views over the quantizer tensors rather than copying them, so they must + * stay open for as long as the index is in use. A dense dataset passed to + * {@link #withDataset(CuVSMatrix)} is owned by the index, as it is for a non-BBQ build, and is + * closed with it. + */ + Builder withBbqDataset(BbqQuantizer... quantizers); + /** * Registers an instance of configured {@link CagraIndexParams} with this * Builder. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index a205931034..491c5213fc 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -50,6 +50,8 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { + + private static final System.Logger LOG = System.getLogger(CagraIndexImpl.class.getName()); private final CuVSResources resources; private final IndexReference cagraIndexReference; private boolean destroyed; @@ -70,6 +72,19 @@ private CagraIndexImpl( this.cagraIndexReference = build(indexParameters, (CuVSMatrixInternal) dataset); } + private CagraIndexImpl( + CagraIndexParams indexParameters, + CuVSMatrix dataset, + BbqQuantizer[] bbqQuantizers, + CuVSResources resources) { + this.resources = resources; + if (dataset != null && !(dataset instanceof CuVSMatrixInternal)) { + throw new IllegalArgumentException("dataset must be a native CuVS matrix"); + } + this.cagraIndexReference = + buildBbq(indexParameters, (CuVSMatrixInternal) dataset, bbqQuantizers); + } + /** * Constructor for loading the index from an {@link InputStream} * @@ -138,14 +153,14 @@ private void checkNotDestroyed() { public void close() { checkNotDestroyed(); try { - int returnValue = cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()); - checkCuVSError(returnValue, "cuvsCagraIndexDestroy"); + // Both are released whatever either reports: `destroyed` below makes a retry impossible. + quietly( + () -> + checkCuVSError( + cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()), + "cuvsCagraIndexDestroy")); if (cagraIndexReference.datasetOwner != null) { - try { - cagraIndexReference.datasetOwner.close(); - } catch (Exception e) { - throw new RuntimeException("Failed to destroy CAGRA dataset", e); - } + quietly(cagraIndexReference.datasetOwner::close); } } finally { destroyed = true; @@ -163,59 +178,224 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInterna long rows = dataset.size(); try (var indexParams = segmentFromIndexParams(indexParameters); - var localArena = Arena.ofConfined()) { + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { MemorySegment indexParamsMemorySegment = indexParams.handle(); int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; - omp_set_num_threads(numWriterThreads); - - var datasetTensor = dataset.toTensor(localArena); + MemorySegment datasetView = MemorySegment.NULL; + boolean success = false; var index = createCagraIndex(); - if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) - == 1) { // when build algo is IVF_PQ - MemorySegment cuvsIvfPqIndexParamsMS = - cuvsIvfPqParams.ivf_pq_build_params( - cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment)); - int n_lists = cuvsIvfPqIndexParams.n_lists(cuvsIvfPqIndexParamsMS); - // As rows cannot be less than n_lists value so trim down. - cuvsIvfPqIndexParams.n_lists( - cuvsIvfPqIndexParamsMS, (int) (rows < n_lists ? rows : n_lists)); - } - try (var resourcesAccessor = resources.access()) { + try { + omp_set_num_threads(numWriterThreads); + + if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) + == 1) { // when build algo is IVF_PQ + MemorySegment cuvsIvfPqIndexParamsMS = + cuvsIvfPqParams.ivf_pq_build_params( + cuvsCagraIndexParams.graph_build_params(indexParamsMemorySegment)); + int n_lists = cuvsIvfPqIndexParams.n_lists(cuvsIvfPqIndexParamsMS); + // As rows cannot be less than n_lists value so trim down. + cuvsIvfPqIndexParams.n_lists( + cuvsIvfPqIndexParamsMS, (int) (rows < n_lists ? rows : n_lists)); + } + var cuvsRes = resourcesAccessor.handle(); // TODO: do we need a stream sync here? - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - MemorySegment datasetView = MemorySegment.NULL; - try { - MemorySegment datasetViewPtr = localArena.allocate(cuvsDataset_t); - if (isCagraPaddedLayout(dataset)) { - returnValue = cuvsDatasetMakePaddedView(cuvsRes, datasetTensor, datasetViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePaddedView"); - } else { - returnValue = cuvsDatasetMakeStandardView(cuvsRes, datasetTensor, datasetViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakeStandardView"); - } - datasetView = datasetViewPtr.get(cuvsDataset_t, 0); + datasetView = + isCagraPaddedLayout(dataset) + ? makePaddedViewHandle(cuvsRes, dataset, localArena) + : makeStandardViewHandle(cuvsRes, dataset, localArena); - returnValue = cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetView, index); - checkCuVSError(returnValue, "cuvsCagraBuild"); - } finally { - if (datasetView.address() != 0) { - checkCuVSError(cuvsDatasetDestroy(datasetView), "cuvsDatasetDestroy"); - } + checkCuVSError( + cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetView, index), + "cuvsCagraBuild"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + + success = true; + return new IndexReference(index, dataset); + } finally { + MemorySegment view = datasetView; + if (view.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(view), "cuvsDatasetDestroy")); + } + // Ownership transfers only on success, so on any other exit the index is ours and the + // caller's matrix is not. + if (!success) { + quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); } + quietly(() -> omp_set_num_threads(1)); + } + } + } - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + private IndexReference buildBbq( + CagraIndexParams indexParameters, CuVSMatrixInternal dataset, BbqQuantizer[] bbqQuantizers) { + if (bbqQuantizers == null || bbqQuantizers.length < 1 || bbqQuantizers.length > 2) { + throw new IllegalArgumentException("BBQ build requires one or two quantizers"); + } + for (BbqQuantizer quantizer : bbqQuantizers) { + Objects.requireNonNull(quantizer); + } + + try (var indexParams = segmentFromIndexParams(indexParameters); + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + var quantizerHandles = new ArrayList(bbqQuantizers.length); + MemorySegment bbqDataset = MemorySegment.NULL; + MemorySegment paddedDataset = MemorySegment.NULL; + int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; + boolean success = false; + var index = createCagraIndex(); + + try { + // nn-descent, which a BBQ build always uses, does a substantial amount of its work on + // the host under OpenMP. + omp_set_num_threads(numWriterThreads); + for (BbqQuantizer quantizer : bbqQuantizers) { + var quantizerPtr = localArena.allocate(cuvsBbqQuantizer_t); + checkCuVSError( + cuvsBbqQuantizerCreateView( + tensor(quantizer.getCodes(), "codes", localArena), + vectorTensor(quantizer.getLowerIntervals(), "lowerIntervals", localArena), + vectorTensor(quantizer.getUpperIntervals(), "upperIntervals", localArena), + vectorTensor( + quantizer.getAdditionalCorrections(), "additionalCorrections", localArena), + vectorTensor( + quantizer.getQuantizedComponentSums(), "quantizedComponentSums", localArena), + vectorTensor(quantizer.getCentroid(), "centroid", localArena), + vectorTensor(quantizer.getDequantDelta(), "dequantDelta", localArena), + vectorTensor(quantizer.getDequantSumDelta(), "dequantSumDelta", localArena), + vectorTensor(quantizer.getRowNorm(), "rowNorm", localArena), + quantizer.getLayout().value, + quantizer.getMetric().value, + quantizer.getCentroidNormSq(), + quantizerPtr), + "cuvsBbqQuantizerCreateView"); + quantizerHandles.add(quantizerPtr.get(cuvsBbqQuantizer_t, 0)); + } + + var quantizerArray = localArena.allocate(ValueLayout.ADDRESS, quantizerHandles.size()); + for (int i = 0; i < quantizerHandles.size(); ++i) { + quantizerArray.setAtIndex(ValueLayout.ADDRESS, i, quantizerHandles.get(i)); + } + var bbqDatasetPtr = localArena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakeBbqView(cuvsRes, quantizerArray, quantizerHandles.size(), bbqDatasetPtr), + "cuvsDatasetMakeBbqView"); + bbqDataset = bbqDatasetPtr.get(cuvsDataset_t, 0); + + checkCuVSError( + cuvsCagraBuild(cuvsRes, indexParams.handle(), bbqDataset, index), "cuvsCagraBuild"); + + IndexReference result; + if (dataset == null) { + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + result = new IndexReference(index, null, null); + } else { + paddedDataset = makeDevicePaddedDataset(cuvsRes, dataset, localArena); + checkCuVSError( + cuvsCagraUpdateDataset(cuvsRes, paddedDataset, index), "cuvsCagraUpdateDataset"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + + // The index owns both the padded dataset and the caller's matrix from here on. + result = + new IndexReference( + index, + dataset, + new PaddedDatasetOwner(new DatasetCloseDelegate(paddedDataset), dataset)); + } + success = true; + return result; + } finally { + // Ownership transfers only on success, so on any other exit the index and the padded + // dataset are ours and the caller's matrix is not. Destroy the index first: it holds a + // view into the padded dataset. + if (!success) { + quietly(() -> checkCuVSError(cuvsCagraIndexDestroy(index), "cuvsCagraIndexDestroy")); + MemorySegment padded = paddedDataset; + if (padded.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(padded), "cuvsDatasetDestroy")); + } + } + MemorySegment bbqView = bbqDataset; + if (bbqView.address() != 0) { + quietly(() -> checkCuVSError(cuvsDatasetDestroy(bbqView), "cuvsDatasetDestroy")); + } + for (MemorySegment quantizer : quantizerHandles) { + quietly( + () -> checkCuVSError(cuvsBbqQuantizerDestroy(quantizer), "cuvsBbqQuantizerDestroy")); + } + quietly(() -> omp_set_num_threads(1)); } - omp_set_num_threads(1); + } + } - return new IndexReference(index, dataset); + /** + * Validates a caller-supplied BBQ component and returns it in its internal form. + * + *

The native side reads every one of these through {@code from_dlpack}, which accepts only + * compact row-major memory, so a padded matrix has to be rejected here. Left to the C layer it + * surfaces as an error naming the quantizer entry point, with no indication of which of the nine + * tensors was at fault. + */ + private static CuVSMatrixInternal bbqComponent(CuVSMatrix matrix, String name) { + if (!(matrix instanceof CuVSDeviceMatrix) || !(matrix instanceof CuVSMatrixInternal internal)) { + throw new IllegalArgumentException("BBQ tensor '" + name + "' must be a device CuVS matrix"); } + long rowStride = internal.rowStride(); + if (rowStride > 0 && rowStride != matrix.columns()) { + throw new IllegalArgumentException( + "BBQ tensor '" + + name + + "' must be contiguous, but its row stride (" + + rowStride + + ") does not match its column count (" + + matrix.columns() + + ")"); + } + return internal; + } + + /** + * The rank-2 codes tensor. + * + *

Built without strides rather than through {@code toTensor}, which emits + * {@code {rowStride, columnStride}} whenever a row stride is set and so would pass a column + * stride of -1 for a matrix that only declared a row stride. Omitting strides is correct here + * because {@link #bbqComponent} has already rejected anything non-contiguous. + */ + private static MemorySegment tensor(CuVSMatrix matrix, String name, Arena arena) { + var internal = bbqComponent(matrix, name); + return prepareTensor( + arena, + internal.memorySegment(), + new long[] {matrix.size(), matrix.columns()}, + internal.code(), + internal.bits(), + kDLCUDA()); + } + + /** + * One of the rank-1 correction vectors, flattened from the caller's matrix. + * + *

Flattening to {@code size * columns} is only valid because {@link #bbqComponent} has + * rejected anything non-contiguous, so there is no row padding to skip. + */ + private static MemorySegment vectorTensor(CuVSMatrix matrix, String name, Arena arena) { + var internal = bbqComponent(matrix, name); + return prepareTensor( + arena, + internal.memorySegment(), + new long[] {matrix.size() * matrix.columns()}, + internal.code(), + internal.bits(), + kDLCUDA()); } private static MemorySegment createCagraIndex() { @@ -422,27 +602,73 @@ public MemorySegment getIndexHandle() { return cagraIndexReference.getMemorySegment(); } - @Override - public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Throwable { - checkNotDestroyed(); + /** Narrows a caller-supplied dataset to the internal form the native layer can read. */ + private static CuVSMatrixInternal requireInternalMatrix(CuVSMatrix dataset) { Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { + if (!(dataset instanceof CuVSMatrixInternal internal)) { throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); } + return internal; + } + + /** Creates a native padded dataset that views the caller's storage rather than copying it. */ + private static MemorySegment makePaddedViewHandle( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakePaddedView(cuvsRes, dataset.toTensor(arena), out), + "cuvsDatasetMakePaddedView"); + return out.get(cuvsDataset_t, 0); + } + + /** Creates a native padded dataset owning its own copy of the caller's storage. */ + private static MemorySegment makePaddedCopyHandle( + long cuvsRes, CuVSMatrixInternal dataset, int targetMemType, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakePadded(cuvsRes, dataset.toTensor(arena), targetMemType, out), + "cuvsDatasetMakePadded"); + return out.get(cuvsDataset_t, 0); + } + + /** Creates a native standard (unpadded) dataset that views the caller's storage. */ + private static MemorySegment makeStandardViewHandle( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + MemorySegment out = arena.allocate(cuvsDataset_t); + checkCuVSError( + cuvsDatasetMakeStandardView(cuvsRes, dataset.toTensor(arena), out), + "cuvsDatasetMakeStandardView"); + return out.get(cuvsDataset_t, 0); + } + + /** + * Produces a device padded dataset for {@code dataset}: a view when it already has CAGRA's padded + * row width, a copy otherwise. Always device-resident, because {@code cuvsCagraUpdateDataset} + * accepts nothing else. + */ + private static MemorySegment makeDevicePaddedDataset( + long cuvsRes, CuVSMatrixInternal dataset, Arena arena) { + return dataset instanceof CuVSDeviceMatrix && isCagraPaddedLayout(dataset) + ? makePaddedViewHandle(cuvsRes, dataset, arena) + : makePaddedCopyHandle(cuvsRes, dataset, CUVS_DATASET_MEM_TYPE_DEVICE(), arena); + } + + @Override + public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Throwable { + checkNotDestroyed(); + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); + // A padded copy stays wherever the source lives; only attaching one to an index forces it + // onto the device. int targetMemType = (datasetInternal instanceof CuVSHostMatrixImpl) ? CUVS_DATASET_MEM_TYPE_HOST() : CUVS_DATASET_MEM_TYPE_DEVICE(); - MemorySegment paddedDatasetPtr = localArena.allocate(cuvsDataset_t); - var returnValue = - cuvsDatasetMakePadded(cuvsRes, datasetTensor, targetMemType, paddedDatasetPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePadded"); - MemorySegment paddedDataset = paddedDatasetPtr.get(cuvsDataset_t, 0); + MemorySegment paddedDataset = + makePaddedCopyHandle(cuvsRes, datasetInternal, targetMemType, localArena); var out = new CagraIndex.PaddedDataset(); out.setDelegate(new DatasetCloseDelegate(paddedDataset), paddedDataset.address()); @@ -453,19 +679,12 @@ public CagraIndex.PaddedDataset makePaddedDataset(CuVSMatrix dataset) throws Thr @Override public CagraIndex.PaddedDatasetView makePaddedDatasetView(CuVSMatrix dataset) throws Throwable { checkNotDestroyed(); - Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { - throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); - } + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); - MemorySegment paddedViewPtr = localArena.allocate(cuvsDataset_t); - var returnValue = cuvsDatasetMakePaddedView(cuvsRes, datasetTensor, paddedViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakePaddedView"); - MemorySegment paddedView = paddedViewPtr.get(cuvsDataset_t, 0); + MemorySegment paddedView = makePaddedViewHandle(cuvsRes, datasetInternal, localArena); var out = new CagraIndex.PaddedDatasetView(); out.setDelegate(new DatasetCloseDelegate(paddedView), paddedView.address()); @@ -477,19 +696,12 @@ public CagraIndex.PaddedDatasetView makePaddedDatasetView(CuVSMatrix dataset) th public CagraIndex.StandardDatasetView makeStandardDatasetView(CuVSMatrix dataset) throws Throwable { checkNotDestroyed(); - Objects.requireNonNull(dataset); - if (!(dataset instanceof CuVSMatrixInternal datasetInternal)) { - throw new IllegalArgumentException("dataset must be a CuVSMatrixInternal matrix"); - } + var datasetInternal = requireInternalMatrix(dataset); try (var localArena = Arena.ofConfined(); var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); - var datasetTensor = datasetInternal.toTensor(localArena); - MemorySegment standardViewPtr = localArena.allocate(cuvsDataset_t); - var returnValue = cuvsDatasetMakeStandardView(cuvsRes, datasetTensor, standardViewPtr); - checkCuVSError(returnValue, "cuvsDatasetMakeStandardView"); - MemorySegment standardView = standardViewPtr.get(cuvsDataset_t, 0); + MemorySegment standardView = makeStandardViewHandle(cuvsRes, datasetInternal, localArena); var out = new CagraIndex.StandardDatasetView(); out.setDelegate(new DatasetCloseDelegate(standardView), standardView.address()); @@ -760,6 +972,44 @@ private IndexReference deserialize( } } + /** A cleanup action that is allowed to fail, including with a checked exception. */ + @FunctionalInterface + private interface CleanupStep { + void run() throws Exception; + } + + /** + * Runs a cleanup step, logging any failure rather than propagating it. + * + * TODO: replace this with a more comprehensive mechanism as part of https://github.com/NVIDIA/cuvs/issues/2670 + */ + private static void quietly(CleanupStep cleanup) { + try { + cleanup.run(); + } catch (Exception thrown) { + // Error is left to propagate: an OutOfMemoryError or a linkage failure says the JVM is in + // no state to carry on, and swallowing it would hide that. + LOG.log( + System.Logger.Level.WARNING, + "Failed to release a native cuVS resource; it has probably been leaked", + thrown); + } + } + + /** + * Owns both halves of an attached dataset: the native {@code cuvsDataset_t} and the matrix whose + * storage it may be viewing. The native dataset is released first, since closing the matrix + * underneath a live view would leave it dangling. + */ + private record PaddedDatasetOwner(DatasetCloseDelegate nativeDataset, CuVSMatrix matrix) + implements AutoCloseable { + @Override + public void close() { + quietly(nativeDataset::close); + quietly(matrix::close); + } + } + private static final class DatasetCloseDelegate implements AutoCloseable { private MemorySegment handle; @@ -1096,6 +1346,7 @@ public static class Builder implements CagraIndex.Builder { private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private CuVSMatrix graph; + private BbqQuantizer[] bbqQuantizers; public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; @@ -1133,6 +1384,12 @@ public Builder withDataset(CuVSMatrix dataset) { return this; } + @Override + public Builder withBbqDataset(BbqQuantizer... quantizers) { + this.bbqQuantizers = quantizers == null ? null : quantizers.clone(); + return this; + } + @Override public Builder withIndexParams(CagraIndexParams cagraIndexParameters) { this.cagraIndexParams = cagraIndexParameters; @@ -1153,6 +1410,8 @@ public CagraIndexImpl build() throws Throwable { } return new CagraIndexImpl( cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources); + } else if (bbqQuantizers != null) { + return new CagraIndexImpl(cagraIndexParams, dataset, bbqQuantizers, cuvsResources); } else if (dataset != null) { return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } else { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java new file mode 100644 index 0000000000..31e659cb96 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBbqIT.java @@ -0,0 +1,673 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import java.lang.invoke.MethodHandles; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * CAGRA indices built from BBQ quantizers via {@link CagraIndex.Builder#withBbqDataset}. + * + *

{@code gpuSearchBaselineWithoutBbq} is a control: the same search over a graph plain CAGRA + * built from the same vectors. If it regresses alongside the BBQ case, look at the search + * parameters or the data before suspecting quantization. + */ +public class CagraBbqIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + private static final int ROWS = 20000; + private static final int DIM = 64; + private static final int QUERIES = 32; + private static final int CLUSTERS = 40; + private static final double CLUSTER_SPREAD = 3.0; + + /** A neighbor within this much of the true k-th distance is as good an answer. */ + private static final float TIE_TOLERANCE = 0.01f; + + private static final int TOP_K = 5; + + /** + * Floor for recall@{@value #TOP_K}. Both searches measure 1.0 on this data with the pinned search + * parameters, so this leaves room for nn-descent's run-to-run variation without being loose + * enough to pass a real regression. + */ + private static final double RECALL_FLOOR = 0.95; + + /** PACKED_1B encodes one bit per dimension. */ + private static final int CODE_BYTES = (DIM + 7) / 8; + + private float[][] dataset; + private float[][] queries; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + // Fixed seed so the dataset and queries are stable. Note this does not make recall stable: + // nn-descent is stochastic, so the measured values still move by a few points between runs. + random = new java.util.Random(42); + float[][] centers = randomVectors(CLUSTERS); + dataset = clusteredVectors(centers, ROWS); + queries = clusteredVectors(centers, QUERIES); + } + + /** Builds the graph from quantizers with a dense dataset attached, then searches on the GPU. */ + @Test + public void gpuBuildGpuSearch() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + // Owned by the index once built, so it is not closed here. + CuVSMatrix dense = newDenseDataset(resources); + try (CagraIndex index = bbqIndex(resources, quantizer, dense); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + assertRecall("gpu-search", index.search(query)); + } + } + } + + /** + * Plain CAGRA over the same vectors, as a reference point for the BBQ recall above. Both sides + * run the same GPU search, so the difference between the two numbers is the cost of building the + * graph from quantized codes. + */ + @Test + public void gpuSearchBaselineWithoutBbq() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create()) { + // Owned by the index once built, so it is not closed here. + CuVSMatrix dense = newDenseDataset(resources); + try (CagraIndex index = + CagraIndex.newBuilder(resources) + .withIndexParams( + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .build()) + .withDataset(dense) + .build(); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + assertRecall("baseline-gpu-search", index.search(query)); + } + } + } + + /** + * Extracting the graph from a graph-only BBQ index, which is how cuvs-lucene consumes one: no + * dense dataset is attached, so this is the memory-efficient shape BBQ exists for, and nothing + * native dereferences the BBQ dataset view on this path. + */ + @Test + public void graphOnlyBbqBuildProducesUsableGraph() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix graph = index.getGraph()) { + assertEquals("one adjacency row per vector", ROWS, graph.size()); + int degree = (int) graph.columns(); + assertTrue("graph degree should be positive", degree > 0); + + int[][] adjacency = readAdjacency(graph); + int selfLoops = 0; + for (int node = 0; node < ROWS; node++) { + for (int neighbor : adjacency[node]) { + assertTrue( + "neighbor id " + neighbor + " of node " + node + " is out of range", + neighbor >= 0 && neighbor < ROWS); + if (neighbor == node) { + selfLoops++; + } + } + } + // A handful of self-references is normal padding; a graph made mostly of them is not. + assertTrue( + "graph is mostly self-references (" + selfLoops + " of " + ROWS * degree + ")", + selfLoops < ROWS * degree / 10); + } + } + } + + /** + * A dense dataset handed to the builder is owned by the resulting index, exactly as it is for a + * non-BBQ build, so closing the index releases it. Host matrices are arena-backed, which makes + * the release observable: reads fail once the arena is gone. + */ + @Test + public void indexOwnsTheDenseDatasetItWasGiven() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + // hostBuilder yields an arena-backed matrix; ofArray's close() is a no-op and would make + // this test pass whether or not the index released anything. + var denseBuilder = CuVSMatrix.hostBuilder(ROWS, DIM, CuVSMatrix.DataType.FLOAT); + for (float[] row : dataset) { + denseBuilder.addVector(row); + } + CuVSMatrix dense = denseBuilder.build(); + try (CagraIndex index = bbqIndex(resources, quantizer, dense)) { + assertNotNull("the dataset should be readable while the index is open", dense.getRow(0)); + } + // The read, not getRow itself, is what validates the arena's scope. + assertThrows( + "closing the index should have released the dataset it was given", + IllegalStateException.class, + () -> dense.getRow(0).getAsFloat(0)); + } + } + + /** + * Every quantizer component is required. {@code centroidNormSq} is the one that cannot rely on a + * null check, since it is a primitive whose unset state would otherwise read as a legitimate + * zero and quietly skew inner-product and cosine distances. + */ + @Test + public void quantizerMustRejectMissingCentroidNormSq() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + assertThrows( + "omitting withCentroidNormSq should fail fast, like every other component does", + RuntimeException.class, + () -> + new BbqQuantizer.Builder() + .withCodes(q.getCodes()) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + // withCentroidNormSq deliberately omitted + .build()); + } + } + + /** A build takes one or two encoded representations; anything else is a usage error. */ + @Test + public void buildRejectsAnUnusableQuantizerCount() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + assertThrows( + "a build needs at least one representation", + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {}, null)); + assertThrows( + "three representations are more than the native side accepts", + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {q, q, q}, null)); + assertThrows( + "a null representation should be caught before it reaches native code", + NullPointerException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {q, null}, null)); + } + } + + /** Quantizer components live on the device; a host matrix cannot be read by the native side. */ + @Test + public void hostQuantizerComponentIsRejected() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var hostBuilder = CuVSMatrix.hostBuilder(ROWS, 1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < ROWS; i++) { + hostBuilder.addVector(new float[] {1.0f}); + } + try (CuVSMatrix hostRowNorm = hostBuilder.build()) { + var withHostComponent = + new BbqQuantizer.Builder() + .withCodes(q.getCodes()) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(hostRowNorm) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + var failure = + assertThrows( + IllegalArgumentException.class, + () -> bbqIndex(resources, new BbqQuantizer[] {withHostComponent}, null)); + assertTrue( + "the message should name the offending component, but was: " + failure.getMessage(), + failure.getMessage().contains("rowNorm")); + } + } + } + + /** A graph-only index carries no vectors, so searching it before attaching a dataset fails. */ + @Test + public void graphOnlyIndexCannotBeSearchedUntilADatasetIsAttached() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + var failure = assertThrows(RuntimeException.class, () -> index.search(query)); + assertTrue( + "the message should point at updateDataset, but was: " + failure.getMessage(), + failure.getMessage().contains("cuvsCagraUpdateDataset")); + } + } + } + + /** The documented follow-up: attach a dataset to a graph-only index and it becomes searchable. */ + @Test + public void attachingADatasetMakesAGraphOnlyIndexSearchable() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources); + CuVSMatrix dense = newDenseDataset(resources)) { + try (CagraIndex index = bbqIndex(resources, quantizer, null); + CuVSMatrix queryVectors = CuVSMatrix.ofArray(queries)) { + try (CagraIndex.PaddedDatasetView padded = index.makePaddedDatasetView(dense)) { + index.updateDataset(padded); + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(TOP_K) + .withSearchParams(searchParams()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + var results = index.search(query).getResults(); + assertEquals("one result set per query", QUERIES, results.size()); + for (var perQuery : results) { + assertEquals("each query should return topK neighbors", TOP_K, perQuery.size()); + } + } + } + } + } + + /** + * A component may declare a row stride explicitly as long as it equals the column count. That is + * still contiguous, so it has to be accepted and passed to the native side without strides. + */ + @Test + public void explicitlyContiguousCodesMatrixIsAccepted() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var codesBuilder = + CuVSMatrix.deviceBuilder( + resources, ROWS, CODE_BYTES, CODE_BYTES, -1, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < ROWS; i++) { + codesBuilder.addVector(new byte[CODE_BYTES]); + } + try (CuVSMatrix stridedCodes = codesBuilder.build()) { + var withStridedCodes = + new BbqQuantizer.Builder() + .withCodes(stridedCodes) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + try (CagraIndex index = bbqIndex(resources, new BbqQuantizer[] {withStridedCodes}, null); + CuVSMatrix graph = index.getGraph()) { + assertEquals("one adjacency row per vector", ROWS, graph.size()); + } + } + } + } + + /** + * Every quantizer component has to be contiguous, because the native side reads them through + * {@code from_dlpack}, which accepts only compact row-major memory. The check belongs on the Java + * side so the message can name the offending tensor rather than just the entry point. + */ + @Test + public void paddedCodesMatrixMustBeRejectedClearly() throws Throwable { + try (CuVSResources resources = CheckedCuVSResources.create(); + Quantizer quantizer = newQuantizer(resources)) { + var q = quantizer.quantizer(); + var paddedCodesBuilder = + CuVSMatrix.deviceBuilder( + resources, ROWS, CODE_BYTES, CODE_BYTES * 2, -1, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < ROWS; i++) { + paddedCodesBuilder.addVector(new byte[CODE_BYTES]); + } + try (CuVSMatrix paddedCodes = paddedCodesBuilder.build()) { + var padded = + new BbqQuantizer.Builder() + .withCodes(paddedCodes) + .withLowerIntervals(q.getLowerIntervals()) + .withUpperIntervals(q.getUpperIntervals()) + .withAdditionalCorrections(q.getAdditionalCorrections()) + .withQuantizedComponentSums(q.getQuantizedComponentSums()) + .withCentroid(q.getCentroid()) + .withDequantDelta(q.getDequantDelta()) + .withDequantSumDelta(q.getDequantSumDelta()) + .withRowNorm(q.getRowNorm()) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(q.getCentroidNormSq()) + .build(); + var failure = + assertThrows( + "a padded codes matrix should be rejected by the Java guard, as vector tensors are", + IllegalArgumentException.class, + () -> + CagraIndex.newBuilder(resources) + .withIndexParams( + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .build()) + .withBbqDataset(padded) + .build()); + assertTrue( + "the message should name the offending component, but was: " + failure.getMessage(), + failure.getMessage().contains("codes")); + } + } + } + + private static int[][] readAdjacency(CuVSMatrix graph) { + int nodes = (int) graph.size(); + int degree = (int) graph.columns(); + int[][] adjacency = new int[nodes][degree]; + for (int i = 0; i < nodes; i++) { + RowView row = graph.getRow(i); + for (int j = 0; j < degree; j++) { + adjacency[i][j] = row.getAsInt(j); + } + } + return adjacency; + } + + /** + * Recall@{@value #TOP_K} against brute-force ground truth, scored by distance rather than by id. + * + *

A returned neighbor counts when it is at least as close as the true {@value #TOP_K}-th, so + * an answer that is equally good but differently tied is not penalised. Scoring by id instead + * measures how the data happens to break ties, which on synthetic vectors is mostly noise. + */ + private void assertRecall(String label, SearchResults results) { + int hits = 0; + int total = 0; + for (int q = 0; q < QUERIES; q++) { + float[] distances = new float[ROWS]; + for (int i = 0; i < ROWS; i++) { + distances[i] = squaredDistance(queries[q], dataset[i]); + } + float[] sorted = distances.clone(); + Arrays.sort(sorted); + float acceptable = sorted[TOP_K - 1] * (1.0f + TIE_TOLERANCE); + for (int id : results.getResults().get(q).keySet()) { + total++; + // CAGRA reports a missing neighbor as an out-of-range id; count it as a miss rather than + // indexing with it. + if (id >= 0 && id < ROWS && distances[id] <= acceptable) { + hits++; + } + } + } + double recall = (double) hits / total; + log.info("{} recall@{} = {}", label, TOP_K, recall); + assertTrue(label + " recall@" + TOP_K + " was only " + recall, recall >= RECALL_FLOOR); + } + + /** + * Search effort is pinned rather than left to the defaults: recall is being asserted here, and a + * default that changes would silently move the number the assertion is checking. + */ + private static CagraSearchParams searchParams() { + return new CagraSearchParams.Builder() + .withAlgo(CagraSearchParams.SearchAlgo.SINGLE_CTA) + .withItopkSize(256) + .build(); + } + + private static float squaredDistance(float[] a, float[] b) { + float sum = 0.0f; + for (int i = 0; i < a.length; i++) { + float diff = a[i] - b[i]; + sum += diff * diff; + } + return sum; + } + + private static CagraIndex bbqIndex( + CuVSResources resources, Quantizer quantizer, CuVSMatrix dataset) throws Throwable { + return bbqIndex(resources, new BbqQuantizer[] {quantizer.quantizer()}, dataset); + } + + private static CagraIndex bbqIndex( + CuVSResources resources, BbqQuantizer[] quantizers, CuVSMatrix dataset) throws Throwable { + var params = + new CagraIndexParams.Builder() + .withMetric(CuvsDistanceType.L2Expanded) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .build(); + var builder = + CagraIndex.newBuilder(resources).withIndexParams(params).withBbqDataset(quantizers); + if (dataset != null) { + builder = builder.withDataset(dataset); + } + return builder.build(); + } + + private float[][] randomVectors(int count) { + float[][] vectors = new float[count][DIM]; + for (float[] vector : vectors) { + for (int d = 0; d < DIM; d++) { + vector[d] = random.nextFloat() * 100.0f; + } + } + return vectors; + } + + /** + * Vectors drawn around a set of cluster centers. + * + *

Uniformly random high-dimensional vectors are close to equidistant from one another, so a + * graph index has nothing to exploit and recall says more about the data than the + * implementation. Clustered data is both more representative and a meaningful thing to assert on. + */ + private float[][] clusteredVectors(float[][] centers, int count) { + float[][] vectors = new float[count][DIM]; + for (float[] vector : vectors) { + float[] center = centers[random.nextInt(centers.length)]; + for (int d = 0; d < DIM; d++) { + vector[d] = center[d] + (float) (random.nextGaussian() * CLUSTER_SPREAD); + } + } + return vectors; + } + + private CuVSMatrix newDenseDataset(CuVSResources resources) { + var builder = CuVSMatrix.deviceBuilder(resources, ROWS, DIM, CuVSMatrix.DataType.FLOAT); + for (float[] row : dataset) { + builder.addVector(row); + } + return builder.build(); + } + + /** A quantizer plus the component matrices it was built from, closed together. */ + private record Quantizer(BbqQuantizer quantizer, List matrices) + implements AutoCloseable { + @Override + public void close() { + matrices.forEach(CuVSMatrix::close); + } + } + + /** + * Encodes {@link #dataset} as PACKED_1B BBQ codes and uploads every component to device memory. + * + *

This mirrors {@code cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp} closely + * enough for the codes to describe the real data, which is what makes the end-to-end searches + * below meaningful. It deliberately skips that file's {@code optimize_intervals} MSE refinement + * and keeps the initial grid estimate: the graph only has to be good, not optimal, because both + * searches score candidates against the dense dataset rather than against the codes. + */ + private Quantizer newQuantizer(CuVSResources resources) { + float[] centroid = new float[DIM]; + for (float[] row : dataset) { + for (int d = 0; d < DIM; d++) { + centroid[d] += row[d]; + } + } + float centroidNormSq = 0.0f; + for (int d = 0; d < DIM; d++) { + centroid[d] /= ROWS; + centroidNormSq += centroid[d] * centroid[d]; + } + + byte[][] codes = new byte[ROWS][CODE_BYTES]; + float[] lower = new float[ROWS]; + float[] upper = new float[ROWS]; + float[] corrections = new float[ROWS]; + int[] sums = new int[ROWS]; + float[] rowNorms = new float[ROWS]; + float[] delta = new float[ROWS]; + float[] sumDelta = new float[ROWS]; + + for (int i = 0; i < ROWS; i++) { + float[] centred = new float[DIM]; + float origNormSq = 0.0f; + float min = Float.MAX_VALUE; + float max = -Float.MAX_VALUE; + float centredNormSq = 0.0f; + double mean = 0.0; + double var = 0.0; + for (int d = 0; d < DIM; d++) { + origNormSq += dataset[i][d] * dataset[i][d]; + centred[d] = dataset[i][d] - centroid[d]; + min = Math.min(min, centred[d]); + max = Math.max(max, centred[d]); + centredNormSq += centred[d] * centred[d]; + double diff = centred[d] - mean; + mean += diff / (d + 1); + var += diff * (centred[d] - mean); + } + double stddev = Math.sqrt(var / DIM); + + // kMinimumMseGrid[0] from the C++ reference, i.e. the 1-bit row. + float a = (float) clamp(-0.798 * stddev + mean, min, max); + float b = (float) clamp(0.798 * stddev + mean, min, max); + // One step for 1-bit codes, so each component lands on 0 or 1. + float step = b - a; + int sum = 0; + for (int d = 0; d < DIM; d++) { + int code = step == 0.0f ? 0 : Math.round((float) (clamp(centred[d], a, b) - a) / step); + sum += code; + codes[i][d / 8] |= (byte) ((code & 1) << (7 - (d % 8))); + } + + lower[i] = a; + upper[i] = b; + // L2Expanded is euclidean, so the correction carries the centred norm. + corrections[i] = centredNormSq; + sums[i] = sum; + rowNorms[i] = origNormSq; + delta[i] = b - a; + sumDelta[i] = delta[i] * sum; + } + + List owned = new ArrayList<>(); + var codesBuilder = + CuVSMatrix.deviceBuilder(resources, ROWS, CODE_BYTES, CuVSMatrix.DataType.BYTE); + for (byte[] row : codes) { + codesBuilder.addVector(row); + } + + BbqQuantizer quantizer = + new BbqQuantizer.Builder() + .withCodes(track(owned, codesBuilder.build())) + .withLowerIntervals(track(owned, deviceVector(resources, lower))) + .withUpperIntervals(track(owned, deviceVector(resources, upper))) + .withAdditionalCorrections(track(owned, deviceVector(resources, corrections))) + .withQuantizedComponentSums(track(owned, deviceVector(resources, sums))) + .withCentroid(track(owned, deviceRow(resources, centroid))) + .withDequantDelta(track(owned, deviceVector(resources, delta))) + .withDequantSumDelta(track(owned, deviceVector(resources, sumDelta))) + .withRowNorm(track(owned, deviceVector(resources, rowNorms))) + .withLayout(BbqQuantizer.CodeLayout.PACKED_1B) + .withMetric(CuvsDistanceType.L2Expanded) + .withCentroidNormSq(centroidNormSq) + .build(); + return new Quantizer(quantizer, owned); + } + + private static double clamp(double value, double min, double max) { + return Math.min(Math.max(value, min), max); + } + + private static CuVSMatrix track(List owned, CuVSMatrix matrix) { + owned.add(matrix); + return matrix; + } + + /** One value per row, which the BBQ view reads as a length-ROWS vector. */ + private static CuVSMatrix deviceVector(CuVSResources resources, float[] values) { + var builder = CuVSMatrix.deviceBuilder(resources, values.length, 1, CuVSMatrix.DataType.FLOAT); + for (float value : values) { + builder.addVector(new float[] {value}); + } + return builder.build(); + } + + private static CuVSMatrix deviceVector(CuVSResources resources, int[] values) { + var builder = CuVSMatrix.deviceBuilder(resources, values.length, 1, CuVSMatrix.DataType.INT); + for (int value : values) { + builder.addVector(new int[] {value}); + } + return builder.build(); + } + + private static CuVSMatrix deviceRow(CuVSResources resources, float[] row) { + var builder = CuVSMatrix.deviceBuilder(resources, 1, row.length, CuVSMatrix.DataType.FLOAT); + builder.addVector(row); + return builder.build(); + } +} From 693f022cffb5616b989937ea1a65725fe021d008 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Tue, 22 Sep 2026 19:06:54 -0400 Subject: [PATCH 11/12] Increase CUDA 12 wheel size (#2676) Due to updates to dependencies and cuvs we have hit our previous limit for CUDA 12 binary size. Bump the limit by 7mb to allow more feature development. Authors: - Robert Maynard (https://github.com/robertmaynard) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) URL: https://github.com/NVIDIA/cuvs/pull/2676 --- ci/validate_wheel.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 30de27cfc8..82a4c7d74d 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -21,7 +21,7 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcuvs" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '375Mi' + --max-allowed-size-compressed '382Mi' ) else PYDISTCHECK_ARGS+=( From 4254ecaa14e183e17388deaf958332f2fb64e21c Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 23 Sep 2026 16:21:13 +0200 Subject: [PATCH 12/12] Update copyright notice --- python/cuvs/cuvs/neighbors/hnsw/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cuvs/cuvs/neighbors/hnsw/__init__.py b/python/cuvs/cuvs/neighbors/hnsw/__init__.py index 3e8541b603..cbac88163f 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/__init__.py +++ b/python/cuvs/cuvs/neighbors/hnsw/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0