diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 69ede8404e..b6a313ecc5 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1398,6 +1398,7 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/composite/index.cu $<$:src/neighbors/cagra.cpp> $<$:src/neighbors/hnsw.cpp> + $<$:src/neighbors/detail/hnsw/external_translate.cu> src/neighbors/ivf_common.cu src/neighbors/ivf_flat_index.cpp ${ivf_flat_build_extend_inst_files} diff --git a/cpp/src/neighbors/detail/cagra/ace_external_plan.hpp b/cpp/src/neighbors/detail/cagra/ace_external_plan.hpp new file mode 100644 index 0000000000..5013394a3e --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/ace_external_plan.hpp @@ -0,0 +1,473 @@ +/* + * 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 + +namespace cuvs::neighbors::cagra::detail { + +inline uint64_t external_checked_add(uint64_t lhs, uint64_t rhs, const char* what) +{ + RAFT_EXPECTS(rhs <= std::numeric_limits::max() - lhs, "overflow computing %s", what); + return lhs + rhs; +} + +inline uint64_t external_checked_mul(uint64_t lhs, uint64_t rhs, const char* what) +{ + RAFT_EXPECTS( + lhs == 0 || rhs <= std::numeric_limits::max() / lhs, "overflow computing %s", what); + return lhs * rhs; +} + +inline uint64_t external_div_rounding_up(uint64_t value, uint64_t divisor) +{ + RAFT_EXPECTS(divisor != 0, "division by zero in external HNSW plan"); + return value / divisor + static_cast(value % divisor != 0); +} + +inline uint64_t external_maximum_partitions(uint64_t rows, uint64_t intermediate_degree) +{ + RAFT_EXPECTS(rows > intermediate_degree, + "external HNSW dataset is too small for one valid CAGRA partition"); + constexpr uint64_t minimum_core_rows = 1000; + uint64_t graph_limit = + external_checked_mul(2, rows, "maximum external partition rows") / + external_checked_add(intermediate_degree, 1, "minimum external partition occurrence count"); + uint64_t core_limit = rows / minimum_core_rows; + return std::max(2, std::min({rows, graph_limit, core_limit})); +} + +struct external_row_range { + uint64_t start = 0; + uint64_t count = 0; +}; + +inline std::vector make_external_sample_ranges(uint64_t rows, + uint64_t sample_rows, + uint64_t max_stripes = 16) +{ + RAFT_EXPECTS(sample_rows > 0 && sample_rows <= rows && max_stripes > 0, + "invalid external HNSW sample range request"); + uint64_t stripe_count = std::min(max_stripes, sample_rows); + if (stripe_count == 1) { return {{(rows - sample_rows) / 2, sample_rows}}; } + uint64_t gap_count = stripe_count - 1; + uint64_t gap_rows = rows - sample_rows; + std::vector ranges; + ranges.reserve(static_cast(stripe_count)); + uint64_t cursor = 0; + for (uint64_t stripe = 0; stripe < stripe_count; ++stripe) { + uint64_t count = + sample_rows / stripe_count + static_cast(stripe < sample_rows % stripe_count); + ranges.push_back({cursor, count}); + cursor = external_checked_add(cursor, count, "external sample range"); + if (stripe + 1 < stripe_count) { + uint64_t gap = gap_rows / gap_count + static_cast(stripe < gap_rows % gap_count); + cursor = external_checked_add(cursor, gap, "external sample gap"); + } + } + RAFT_EXPECTS(cursor == rows, "external HNSW sample ranges do not cover the requested span"); + return ranges; +} + +inline std::vector make_external_monotonic_ranges(uint64_t rows, + uint64_t chunk_rows) +{ + RAFT_EXPECTS(rows > 0 && chunk_rows > 0, "invalid external HNSW monotonic range request"); + std::vector ranges; + ranges.reserve(static_cast(external_div_rounding_up(rows, chunk_rows))); + for (uint64_t start = 0; start < rows;) { + uint64_t count = std::min(chunk_rows, rows - start); + ranges.push_back({start, count}); + start = external_checked_add(start, count, "external monotonic range"); + } + return ranges; +} + +struct ace_external_plan_input { + uint64_t rows = 0; + uint64_t dim = 0; + uint64_t element_size = 0; + uint64_t index_size = sizeof(uint32_t); + uint64_t M = 0; + uint64_t intermediate_degree = 0; + uint64_t graph_degree = 0; + uint64_t requested_partitions = 0; + uint64_t available_host_bytes = 0; + uint64_t available_device_bytes = 0; + uint64_t optimize_host_fixed = 0; + uint64_t optimize_device_fixed = 0; + uint64_t optimize_host_per_row = 0; + uint64_t optimize_device_per_row = 0; + bool force_disk = false; + bool hierarchy = true; + uint32_t requested_queue_depth = 2; +}; + +struct ace_external_byte_ledger { + uint64_t source_scan = 0; + uint64_t centroid_sample = 0; + uint64_t stage_write = 0; + uint64_t stage_read = 0; + uint64_t base_output = 0; + uint64_t upper_sidecar_write = 0; + uint64_t upper_sidecar_read = 0; + uint64_t final_upper_output = 0; + uint64_t expected_upper_occurrences = 0; + + [[nodiscard]] uint64_t logical_total() const + { + uint64_t total = external_checked_add(source_scan, centroid_sample, "source and sample bytes"); + total = external_checked_add(total, stage_write, "logical byte total"); + total = external_checked_add(total, stage_read, "logical byte total"); + total = external_checked_add(total, base_output, "logical byte total"); + total = external_checked_add(total, upper_sidecar_write, "logical byte total"); + total = external_checked_add(total, upper_sidecar_read, "logical byte total"); + return external_checked_add(total, final_upper_output, "logical byte total"); + } +}; + +struct ace_external_plan { + bool use_disk = false; + uint64_t partitions = 0; + uint64_t target_occurrences = 0; + uint64_t max_occurrences = 0; + uint64_t assignment_chunk_rows = 0; + uint64_t centroid_sample_rows = 0; + uint64_t staging_buffer_bytes = 0; + uint64_t preferred_buffer_bytes = 0; + uint64_t hnsw_output_buffer_bytes = 0; + uint64_t global_upper_level_max_rows = 0; + uint32_t queue_depth = 1; + uint64_t host_budget_bytes = 0; + uint64_t device_budget_bytes = 0; + uint64_t host_peak_bytes = 0; + uint64_t device_peak_bytes = 0; + uint64_t centroid_host_peak_bytes = 0; + uint64_t centroid_device_peak_bytes = 0; + uint64_t assignment_host_peak_bytes = 0; + uint64_t assignment_device_peak_bytes = 0; + uint64_t partition_host_peak_bytes = 0; + uint64_t partition_device_peak_bytes = 0; + uint64_t hierarchy_host_peak_bytes = 0; + uint64_t hierarchy_device_peak_bytes = 0; + uint64_t host_fixed_bytes = 0; + uint64_t device_fixed_bytes = 0; + uint64_t host_per_occurrence = 0; + uint64_t host_reader_per_occurrence = 0; + uint64_t device_per_occurrence = 0; + ace_external_byte_ledger bytes; +}; + +inline uint64_t estimate_materialized_cagra_ace_host_bytes(const ace_external_plan_input& in, + uint64_t partitions) +{ + RAFT_EXPECTS(partitions > 0, "CAGRA-ACE host estimate requires at least one partition"); + const uint64_t vector_bytes = + external_checked_mul(in.dim, in.element_size, "CAGRA-ACE vector row"); + const uint64_t mapping_bytes = external_checked_mul( + in.rows, + external_checked_mul(4, in.index_size, "CAGRA-ACE labels and mappings row"), + "CAGRA-ACE labels and mappings"); + const uint64_t full_graph_bytes = external_checked_mul( + in.rows, + external_checked_mul(in.graph_degree, in.index_size, "CAGRA-ACE final graph row"), + "CAGRA-ACE final graph"); + const uint64_t max_occurrences = external_checked_mul( + 6, external_div_rounding_up(in.rows, partitions), "CAGRA-ACE skewed partition rows"); + const uint64_t partition_graph_bytes = external_checked_mul( + external_checked_add(in.intermediate_degree, in.graph_degree, "CAGRA-ACE combined degree"), + in.index_size, + "CAGRA-ACE partition graph row"); + const uint64_t partition_bytes = external_checked_mul( + max_occurrences, + external_checked_add( + external_checked_add(vector_bytes, partition_graph_bytes, "CAGRA-ACE partition row"), + in.optimize_host_per_row, + "CAGRA-ACE partition row with optimization"), + "CAGRA-ACE maximum partition"); + uint64_t total = + external_checked_add(mapping_bytes, full_graph_bytes, "CAGRA-ACE materialized graph"); + total = external_checked_add(total, in.optimize_host_fixed, "CAGRA-ACE fixed workspace"); + return external_checked_add(total, partition_bytes, "CAGRA-ACE materialized host bytes"); +} + +inline ace_external_byte_ledger make_external_byte_ledger(const ace_external_plan_input& in, + uint64_t centroid_sample_rows) +{ + const uint64_t vector_bytes = external_checked_mul(in.dim, in.element_size, "vector byte size"); + const uint64_t dataset_bytes = external_checked_mul(in.rows, vector_bytes, "dataset byte size"); + const uint64_t sample_bytes = + external_checked_mul(centroid_sample_rows, vector_bytes, "centroid sample bytes"); + const uint64_t stage_record_bytes = + external_checked_add(vector_bytes, 2 * sizeof(uint32_t), "stage record byte size"); + const uint64_t stage_one_direction = external_checked_mul( + external_checked_mul(in.rows, 2, "stage occurrence count"), stage_record_bytes, "stage bytes"); + const uint64_t base_row_bytes = external_checked_add( + external_checked_add( + sizeof(uint32_t), + external_checked_mul(in.graph_degree, in.index_size, "base neighbor bytes"), + "base links"), + external_checked_add(vector_bytes, sizeof(size_t), "base vector and label"), + "base row bytes"); + const uint64_t base_output = + external_checked_mul(in.rows, base_row_bytes, "base HNSW output bytes"); + + const uint64_t expected_upper_occurrences = in.hierarchy && in.M > 1 ? in.rows / (in.M - 1) : 0; + const uint64_t sidecar_record_bytes = + external_checked_add(sizeof(uint32_t), + external_checked_mul(in.M, in.index_size, "upper sidecar links"), + "upper sidecar record"); + const uint64_t sidecar_one_direction = + external_checked_mul(expected_upper_occurrences, sidecar_record_bytes, "upper sidecar bytes"); + const uint64_t upper_node_headers = + external_checked_mul(in.rows, sizeof(uint32_t), "upper node headers"); + const uint64_t upper_block_bytes = external_checked_add( + sizeof(uint32_t), external_checked_mul(in.M, in.index_size, "upper links"), "upper block"); + const uint64_t upper_links = + external_checked_mul(expected_upper_occurrences, upper_block_bytes, "upper output links"); + + return {dataset_bytes, + sample_bytes, + stage_one_direction, + stage_one_direction, + base_output, + sidecar_one_direction, + sidecar_one_direction, + external_checked_add(upper_node_headers, upper_links, "final upper output"), + expected_upper_occurrences}; +} + +inline ace_external_plan make_ace_external_plan(const ace_external_plan_input& in) +{ + RAFT_EXPECTS(in.rows > 0 && in.dim > 0 && in.element_size > 0, + "external HNSW plan requires a non-empty dataset"); + RAFT_EXPECTS(in.rows <= std::numeric_limits::max(), + "external HNSW build supports at most UINT32_MAX rows"); + RAFT_EXPECTS(in.requested_partitions <= in.rows, + "ACE: number of partitions cannot exceed dataset size"); + RAFT_EXPECTS(in.dim <= std::numeric_limits::max(), + "external HNSW build supports at most UINT32_MAX dimensions"); + RAFT_EXPECTS(in.M >= 2 && in.graph_degree > 0 && in.intermediate_degree >= in.graph_degree, + "invalid graph degrees in external HNSW plan"); + + constexpr uint64_t one_mib = uint64_t{1} << 20; + constexpr uint64_t one_gib = uint64_t{1} << 30; + const uint64_t host_limit = + in.available_host_bytes == 0 ? std::numeric_limits::max() : in.available_host_bytes; + const uint64_t device_limit = in.available_device_bytes == 0 + ? std::numeric_limits::max() + : in.available_device_bytes; + const uint64_t host_budget = host_limit - host_limit / 5; + const uint64_t device_budget = device_limit - device_limit / 5; + const uint64_t vector_bytes = external_checked_mul(in.dim, in.element_size, "vector byte size"); + + ace_external_plan out; + out.host_budget_bytes = host_budget; + out.device_budget_bytes = device_budget; + uint64_t maximum_partitions = external_maximum_partitions(in.rows, in.intermediate_degree); + out.partitions = std::max(2, in.requested_partitions); + out.partitions = std::min(out.partitions, maximum_partitions); + + // Cap ACE's 1% centroid sample at 8 GiB and 25% of each memory budget. + uint64_t sample_cap = 8 * one_gib; + if (host_budget != std::numeric_limits::max()) { + sample_cap = std::min(sample_cap, host_budget / 4); + } + if (device_budget != std::numeric_limits::max()) { + sample_cap = std::min(sample_cap, device_budget / 4); + } + // Core plus spill, with 3x imbalance: six occurrences per row / partition count. + out.host_per_occurrence = external_checked_add( + external_checked_add( + vector_bytes, + external_checked_mul(in.graph_degree, in.index_size, "poststage graph row"), + "resident partition row"), + external_checked_add( + in.optimize_host_per_row, + external_checked_add(external_checked_mul(2, in.index_size, "resident partition mappings"), + in.hierarchy ? sizeof(uint8_t) : 0, + "resident partition mappings and hierarchy level"), + "host optimization and resident partition metadata row"), + "host bytes per partition occurrence"); + out.host_reader_per_occurrence = external_checked_add( + vector_bytes, + external_checked_mul(2, in.index_size, "reader mapping and core label row"), + "host bytes per prefetched partition occurrence"); + out.device_per_occurrence = external_checked_add( + external_checked_add( + vector_bytes, + external_checked_mul( + external_checked_add(in.intermediate_degree, in.graph_degree, "combined graph degree"), + in.index_size, + "device graph rows"), + "device partition row"), + external_checked_add( + in.index_size, in.optimize_device_per_row, "device mapping and optimize row"), + "device bytes per partition occurrence"); + + auto update_peaks = [&] { + uint64_t requested_sample = std::max( + external_checked_mul(100, out.partitions, "minimum centroid samples"), in.rows / 100); + uint64_t sample_rows = std::min(requested_sample, in.rows); + sample_rows = std::min(sample_rows, std::max(1, sample_cap / sizeof(float) / in.dim)); + RAFT_EXPECTS(sample_rows >= out.partitions, + "host/device cap is too small for one centroid sample per partition"); + out.centroid_sample_rows = sample_rows; + out.host_fixed_bytes = in.optimize_host_fixed; + out.device_fixed_bytes = in.optimize_device_fixed; + + out.target_occurrences = external_div_rounding_up( + external_checked_mul(2, in.rows, "target partition occurrences"), out.partitions); + out.max_occurrences = external_div_rounding_up( + external_checked_mul(in.rows, 6, "skewed partition rows"), out.partitions); + out.host_peak_bytes = external_checked_add( + out.host_fixed_bytes, + external_checked_mul( + out.max_occurrences, out.host_per_occurrence, "external host partition peak"), + "external host peak"); + out.device_peak_bytes = external_checked_add( + out.device_fixed_bytes, + external_checked_mul( + out.max_occurrences, out.device_per_occurrence, "external device partition peak"), + "external device peak"); + }; + update_peaks(); + + while ((out.host_peak_bytes > host_budget || out.device_peak_bytes > device_budget) && + out.partitions < maximum_partitions) { + uint64_t next = + std::min(maximum_partitions, + std::max(external_checked_add(out.partitions, 1, "next partition count"), + external_checked_mul(out.partitions, 2, "next partition count"))); + out.partitions = next; + update_peaks(); + } + RAFT_EXPECTS(out.host_peak_bytes <= host_budget, + "external HNSW host cap is below the minimum planned partition peak"); + RAFT_EXPECTS(out.device_peak_bytes <= device_budget, + "external HNSW device cap is below the minimum planned partition peak"); + + uint64_t host_headroom = host_budget - out.host_peak_bytes; + RAFT_EXPECTS(host_headroom >= 2 * one_mib, + "external HNSW host cap leaves less than 2 MiB for staging and output buffers"); + const uint64_t second_partition_host_bytes = external_checked_mul( + out.max_occurrences, out.host_reader_per_occurrence, "prefetched partition host peak"); + out.queue_depth = 1; + if (in.requested_queue_depth > 1 && second_partition_host_bytes <= host_headroom - 2 * one_mib) { + out.queue_depth = 2; + host_headroom -= second_partition_host_bytes; + } + constexpr uint64_t minimum_buffer = uint64_t{64} << 10; + uint64_t minimum_stage_buffer = std::max( + one_mib, external_checked_add(vector_bytes, 2 * sizeof(uint32_t), "minimum stage buffer")); + uint64_t minimum_preferred_buffer = std::max( + minimum_buffer, + external_checked_add(vector_bytes, 2 * sizeof(uint32_t), "minimum preferred buffer")); + uint64_t minimum_output_buffer = std::max( + minimum_buffer, + external_checked_mul(in.graph_degree, in.index_size, "minimum graph output buffer")); + RAFT_EXPECTS(minimum_stage_buffer <= host_budget, + "external HNSW host cap is too small for one stage record"); + out.staging_buffer_bytes = + std::max(minimum_stage_buffer, std::min(host_budget / 4, 256 * one_mib)); + out.preferred_buffer_bytes = + std::max(minimum_preferred_buffer, std::min(host_headroom / 8, 64 * one_mib)); + uint64_t concurrent_reader_buffer = out.queue_depth > 1 ? out.preferred_buffer_bytes : 0; + uint64_t two_output_buffers = + external_checked_mul(2, minimum_output_buffer, "minimum output buffers"); + RAFT_EXPECTS(two_output_buffers <= host_headroom && + concurrent_reader_buffer <= host_headroom - two_output_buffers, + "external HNSW host cap is too small for bounded I/O buffers"); + out.hnsw_output_buffer_bytes = + std::max(minimum_output_buffer, + std::min((host_headroom - concurrent_reader_buffer) / 2, 64 * one_mib)); + + uint64_t centroid_bytes = external_checked_mul( + out.partitions, external_checked_mul(in.dim, sizeof(float), "centroid row"), "centroid bytes"); + uint64_t assignment_host_row = + external_checked_add(external_checked_mul(in.dim, sizeof(float), "assignment host input row"), + 2 * sizeof(uint32_t), + "assignment host row"); + uint64_t assignment_device_row = external_checked_add( + external_checked_mul(in.dim, sizeof(float), "assignment device input row"), + external_checked_add( + external_checked_mul(out.partitions, sizeof(float), "assignment distance row"), + 2 * sizeof(float) + 2 * sizeof(uint32_t), + "assignment device output row"), + "assignment device row"); + RAFT_EXPECTS(out.staging_buffer_bytes < host_budget && centroid_bytes < device_budget, + "external HNSW memory cap is too small for assignment buffers"); + uint64_t assignment_host_rows = (host_budget - out.staging_buffer_bytes) / assignment_host_row; + uint64_t assignment_device_rows = (device_budget - centroid_bytes) / assignment_device_row; + out.assignment_chunk_rows = + std::min({uint64_t{32} * 1024, assignment_host_rows, assignment_device_rows}); + RAFT_EXPECTS(out.assignment_chunk_rows > 0, + "external HNSW memory cap cannot hold one assignment row"); + + uint64_t sample_float_bytes = + external_checked_mul(out.centroid_sample_rows, + external_checked_mul(in.dim, sizeof(float), "centroid sample row"), + "centroid sample bytes"); + out.centroid_host_peak_bytes = sample_float_bytes; + out.centroid_device_peak_bytes = + external_checked_add(sample_float_bytes, centroid_bytes, "centroid training device peak"); + out.assignment_host_peak_bytes = external_checked_add( + out.staging_buffer_bytes, + external_checked_mul(out.assignment_chunk_rows, assignment_host_row, "assignment host chunk"), + "assignment host peak"); + out.assignment_device_peak_bytes = external_checked_add( + centroid_bytes, + external_checked_mul( + out.assignment_chunk_rows, assignment_device_row, "assignment device chunk"), + "assignment device peak"); + out.partition_host_peak_bytes = external_checked_add( + out.host_peak_bytes, + external_checked_add( + out.queue_depth > 1 ? second_partition_host_bytes : 0, + external_checked_add( + concurrent_reader_buffer, + external_checked_mul(2, out.hnsw_output_buffer_bytes, "external output and graph buffers"), + "external partition buffers"), + "external queued partition buffers"), + "external partition host peak"); + out.partition_device_peak_bytes = external_checked_add( + out.device_peak_bytes, out.hnsw_output_buffer_bytes, "external translated graph chunk"); + + out.global_upper_level_max_rows = std::max( + 1, + std::min(host_budget / 4 / std::max(1, out.host_per_occurrence), + device_budget / 4 / std::max(1, out.device_per_occurrence))); + out.hierarchy_host_peak_bytes = external_checked_add( + external_checked_mul( + out.global_upper_level_max_rows, out.host_per_occurrence, "external hierarchy host rows"), + external_checked_add( + out.preferred_buffer_bytes, out.hnsw_output_buffer_bytes, "external hierarchy host buffers"), + "external hierarchy host peak"); + out.hierarchy_device_peak_bytes = external_checked_mul( + out.global_upper_level_max_rows, out.device_per_occurrence, "external hierarchy device peak"); + out.host_peak_bytes = std::max({out.centroid_host_peak_bytes, + out.assignment_host_peak_bytes, + out.partition_host_peak_bytes, + out.hierarchy_host_peak_bytes}); + out.device_peak_bytes = std::max({out.centroid_device_peak_bytes, + out.assignment_device_peak_bytes, + out.partition_device_peak_bytes, + out.hierarchy_device_peak_bytes}); + RAFT_EXPECTS(out.host_peak_bytes <= host_budget && out.device_peak_bytes <= device_budget, + "external HNSW phase buffers exceed the configured memory cap"); + + out.bytes = make_external_byte_ledger(in, out.centroid_sample_rows); + + out.use_disk = in.force_disk; + return out; +} + +} // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index bac127294e..74141de4e7 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -42,10 +43,25 @@ #include #include #include +#include #include namespace cuvs::neighbors::hnsw::detail { +template +void all_neighbors_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view neighbors, + cuvs::distance::DistanceType metric); + +} // namespace cuvs::neighbors::hnsw::detail + +#include "hnsw/external_build.cuh" +#include "hnsw/index_impl.hpp" +#include "hnsw/serialize_layout.hpp" + +namespace cuvs::neighbors::hnsw::detail { + class exclusive_hnsw_output_file { public: explicit exclusive_hnsw_output_file(std::filesystem::path output_path) @@ -131,36 +147,6 @@ inline constexpr bool is_device_cagra_hnsw_export_index_v = std::is_same_v> || std::is_same_v>; -// This is needed as hnswlib hardcodes the distance type to float -// or int32_t in certain places. However, we can solve uint8 or int8 -// natively with the patch cuVS applies. We could potentially remove -// all the hardcodes and propagate templates throughout hnswlib, but -// as of now it's not needed. -template -struct hnsw_dist_t { - using type = void; -}; - -template <> -struct hnsw_dist_t { - using type = float; -}; - -template <> -struct hnsw_dist_t { - using type = float; -}; - -template <> -struct hnsw_dist_t { - using type = int; -}; - -template <> -struct hnsw_dist_t { - using type = int; -}; - // Map the dataset element type to a cudaDataType_t. This is a host-only helper that // intentionally avoids pulling CUDA/device dependencies. template @@ -186,125 +172,6 @@ inline cudaDataType_t to_cuda_data_type() return CUDA_R_8U; } -template -struct index_impl : index { - public: - /** - * @brief load a base-layer-only hnswlib index originally saved from a built CAGRA index - * - * @param[in] filepath path to the index - * @param[in] dim dimensions of the training dataset - * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") - * @param[in] hierarchy hierarchy used for upper HNSW layers - */ - index_impl(int dim, cuvs::distance::DistanceType metric, HnswHierarchy hierarchy) - : index{dim, metric, hierarchy} - { - if (metric == cuvs::distance::DistanceType::InnerProduct) { - space_ = std::make_unique::type>>(dim); - } else if (metric == cuvs::distance::DistanceType::L2Expanded) { - if constexpr (std::is_same_v || std::is_same_v) { - space_ = std::make_unique::type>>(dim); - } else if constexpr (std::is_same_v or std::is_same_v) { - space_ = std::make_unique>(dim); - } - } - - RAFT_EXPECTS(space_ != nullptr, "Unsupported metric type was used"); - } - - /** - @brief Get hnswlib index - */ - auto get_index() const -> void const* override { return appr_alg_.get(); } - - /** - @brief Set ef for search - */ - void set_ef(int ef) const override - { - ensure_loaded(); - appr_alg_->ef_ = ef; - } - - /** - @brief Set index - */ - void set_index(std::unique_ptr::type>>&& index) - { - appr_alg_ = std::move(index); - } - - /** - @brief Get space - */ - auto get_space() const -> hnswlib::SpaceInterface::type>* - { - return space_.get(); - } - - /** - @brief Set file descriptor for disk-backed index - */ - void set_file_descriptor(cuvs::util::file_descriptor&& fd) { hnsw_fd_.emplace(std::move(fd)); } - - /** - @brief Get file descriptor - */ - auto file_descriptor() const -> const std::optional& - { - return hnsw_fd_; - } - - /** - @brief Get file path for disk-backed index - */ - std::string file_path() const override - { - if (hnsw_fd_.has_value() && hnsw_fd_->is_valid()) { return hnsw_fd_->get_path(); } - return ""; - } - - /** - @brief Ensure the index is loaded into memory. - If the index is disk-backed and not yet loaded, this will load it from the file. - */ - void ensure_loaded() const - { - if (appr_alg_ != nullptr) { return; } // Already loaded - - // Check if we have a file descriptor to load from - if (!hnsw_fd_.has_value() || !hnsw_fd_->is_valid()) { - RAFT_FAIL("Cannot load HNSW index: no file descriptor available and index not in memory"); - } - - std::string filepath = hnsw_fd_->get_path(); - RAFT_EXPECTS(!filepath.empty(), "Cannot load HNSW index: file path is empty"); - RAFT_EXPECTS(std::filesystem::exists(filepath), - "Cannot load HNSW index: file does not exist: %s", - filepath.c_str()); - - RAFT_LOG_INFO("Loading HNSW index from disk: %s", filepath.c_str()); - - try { - appr_alg_ = std::make_unique::type>>( - space_.get(), filepath); - if (this->hierarchy() == HnswHierarchy::NONE) { appr_alg_->base_layer_only = true; } - } catch (const std::bad_alloc& e) { - RAFT_FAIL( - "Failed to load HNSW index from '%s': insufficient host memory. " - "The index is too large to fit in available RAM. " - "Consider using a machine with more memory or reducing the dataset size.", - filepath.c_str()); - } - } - - private: - mutable std::unique_ptr::type>> appr_alg_; - std::unique_ptr::type>> space_; - std::optional hnsw_fd_; -}; - template std::enable_if_t, std::unique_ptr>> @@ -516,8 +383,7 @@ void serialize_to_hnswlib_batched(raft::resources const& res, // initialize dummy HNSW index to retrieve constants auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); - int odd_graph_degree = graph_degree_int % 2; - auto appr_algo = std::make_unique::type>>( + auto appr_algo = std::make_unique::type>>( hnsw_index->get_space(), 1, (graph_degree_int + 1) / 2, params.ef_construction); bool create_hierarchy = params.hierarchy != HnswHierarchy::NONE; @@ -559,6 +425,8 @@ void serialize_to_hnswlib_batched(raft::resources const& res, // 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; + auto serialize_layout = + external::hnsw_serialize_layout_from_algorithm(*appr_algo, n_rows, dim, graph_degree_int); // write header information RAFT_LOG_DEBUG("Writing HNSW header: offsetLevel0=%zu, n_rows=%zu, size_data_per_element=%zu", @@ -572,33 +440,7 @@ void serialize_to_hnswlib_batched(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)); + external::write_hnsw_header_fields(os, serialize_layout); // host queries auto host_query_set = @@ -615,7 +457,6 @@ void serialize_to_hnswlib_batched(raft::resources const& res, RAFT_LOG_INFO("Writing base level"); size_t bytes_written = 0; float GiB = 1 << 30; - IdxT zero = 0; RAFT_EXPECTS(appr_algo->size_data_per_element_ == dim * sizeof(T) + appr_algo->maxM0_ * sizeof(IdxT) + sizeof(int) + sizeof(size_t), "Size data per element mismatch"); @@ -636,31 +477,17 @@ void serialize_to_hnswlib_batched(raft::resources const& res, for (int64_t batch_idx = 0; batch_idx < current_batch_size; batch_idx++) { const int64_t i = batch_start + batch_idx; - os.write(reinterpret_cast(&graph_degree_int), sizeof(int)); - const IdxT* graph_row = &graph_buffer(batch_idx, 0); - os.write(reinterpret_cast(graph_row), sizeof(IdxT) * graph_degree_int); - - if (odd_graph_degree) { - RAFT_EXPECTS(odd_graph_degree == static_cast(appr_algo->maxM0_) - graph_degree_int, - "Odd graph degree mismatch"); - os.write(reinterpret_cast(&zero), sizeof(IdxT)); - } - - const T* data_row = &dataset_buffer(batch_idx, 0); - os.write(reinterpret_cast(data_row), sizeof(T) * dim); + const T* data_row = &dataset_buffer(batch_idx, 0); + static_assert(std::is_same_v, + "hnswlib serialization requires uint32_t neighbor IDs"); + external::write_hnsw_base_row( + os, serialize_layout, graph_row, data_row, label_buffer(batch_idx)); if (create_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))); + std::memcpy(&host_query_set(order_bw[i] - hist[0], 0), data_row, dim * sizeof(T)); } - - // assign original label - auto label = static_cast(label_buffer(batch_idx)); - os.write(reinterpret_cast(&label), sizeof(std::size_t)); - bytes_written += appr_algo->size_data_per_element_; const auto end_clock = std::chrono::system_clock::now(); @@ -716,11 +543,12 @@ void serialize_to_hnswlib_batched(raft::resources const& res, bytes_written = 0; start_clock = std::chrono::system_clock::now(); + std::vector converted_neighbors(appr_algo->M_); for (int64_t i = 0; i < n_rows; i++) { size_t cur_level = create_hierarchy ? levels[i] : 0; + external::write_hnsw_upper_node_header(os, serialize_layout, cur_level); unsigned int linkListSize = create_hierarchy && cur_level > 0 ? appr_algo->size_links_per_element_ * cur_level : 0; - os.write(reinterpret_cast(&linkListSize), sizeof(int)); bytes_written += sizeof(int); if (linkListSize) { for (size_t pt_level = 1; pt_level <= cur_level; pt_level++) { @@ -729,15 +557,11 @@ void serialize_to_hnswlib_batched(raft::resources const& res, 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)); + converted_neighbors[j] = order[neighbors[j] + offsets[pt_level - 1]]; } + external::write_hnsw_upper_block(os, serialize_layout, converted_neighbors.data(), extent); auto remainder = appr_algo->M_ - neighbor_view.extent(1); - for (size_t j = 0; j < remainder; j++) { - os.write(reinterpret_cast(&zero), sizeof(IdxT)); - } bytes_written += (neighbor_view.extent(1) + remainder) * sizeof(IdxT) + sizeof(int); RAFT_EXPECTS(appr_algo->size_links_per_element_ == (neighbor_view.extent(1) + remainder) * sizeof(IdxT) + sizeof(int), @@ -1264,9 +1088,8 @@ from_cagra(raft::resources const& res, #pragma omp parallel for num_threads(num_threads) for (auto i = start_idx; i < end_idx; i++) { auto pt_id = order[i]; - std::copy(appr_algo->getDataByInternalId(pt_id), - appr_algo->getDataByInternalId(pt_id) + dim, - reinterpret_cast(&host_query_set(i - start_idx, 0))); + std::memcpy( + &host_query_set(i - start_idx, 0), appr_algo->getDataByInternalId(pt_id), dim * sizeof(T)); } // find neighbors of the query set @@ -1338,6 +1161,39 @@ from_cagra(raft::resources const& res, return hnsw_index; } +template +size_t estimate_hnsw_host_memory(int64_t n_rows, + int64_t dim, + int graph_degree, + cuvs::distance::DistanceType metric, + HnswHierarchy hierarchy, + int ef_construction) +{ + RAFT_EXPECTS(n_rows > 0 && dim > 0 && graph_degree > 0, + "HNSW host-memory estimate requires a positive shape and graph degree"); + + auto dummy_index = std::make_unique>(dim, metric, hierarchy); + auto dummy_algo = std::make_unique::type>>( + dummy_index->get_space(), 1, (graph_degree + 1) / 2, ef_construction); + + size_t per_element = dummy_algo->size_data_per_element_; + per_element += sizeof(void*) + sizeof(int) + sizeof(std::mutex); + per_element += 56; // unordered_map node + bucket slot (upper bound) + if (hierarchy != HnswHierarchy::NONE) { + int m_used = std::max(2, (graph_degree + 1) / 2); + size_t size_links_per_element = + static_cast(m_used) * sizeof(uint32_t) + sizeof(uint32_t); + per_element += + static_cast(size_links_per_element / std::log(static_cast(m_used))); + } + + const size_t rows = static_cast(n_rows); + if (rows > std::numeric_limits::max() / per_element) { + return std::numeric_limits::max(); + } + return rows * per_element; +} + inline std::pair get_available_memory( std::optional max_host_memory_gb = std::nullopt, std::optional max_gpu_memory_gb = std::nullopt) @@ -1426,33 +1282,14 @@ std::unique_ptr> from_cagra( int64_t dim = dataset.has_value() ? dataset->extent(1) : cagra_index.dim(); int graph_degree_int = static_cast(cagra_index.graph().extent(1)); - // Instantiate a size-1 dummy to read the exact per-element host footprint. - auto dummy_index = std::make_unique>(dim, cagra_index.metric(), params.hierarchy); - auto dummy_algo = std::make_unique::type>>( - dummy_index->get_space(), 1, (graph_degree_int + 1) / 2, params.ef_construction); - - // The contiguous level-0 array (size_data_per_element_ = level-0 links + vector + label) - // dominates, but hnswlib allocates several additional per-element structures that are NOT - // part of it. These fixed-size costs (locks, hash map) don't shrink with the vector, so a - // flat percentage under-counts for low-dim/8-bit/small-M and hierarchical indexes. Model - // them explicitly instead: - // - linkLists_ : char* per element - // - element_levels_ : int per element - // - link_list_locks_ : std::mutex per element (kept for the index lifetime) - // - label_lookup_ : unordered_map node + bucket slot (~56 B upper bound) - // - upper-level link lists (hierarchy != NONE only): expected - // size_links_per_element / ln(M) bytes per element - size_t per_element = dummy_algo->size_data_per_element_; - per_element += sizeof(void*) + sizeof(int) + sizeof(std::mutex); - per_element += 56; // unordered_map node + bucket slot (upper bound) - if (params.hierarchy != HnswHierarchy::NONE) { - int m_used = std::max(2, (graph_degree_int + 1) / 2); - size_t size_links_per_element = - static_cast(m_used) * sizeof(uint32_t) + sizeof(uint32_t); - per_element += - static_cast(size_links_per_element / std::log(static_cast(m_used))); - } - size_t required_host = static_cast(n_rows) * per_element; + // Account for the contiguous level-0 storage plus hnswlib's per-element pointers, levels, + // locks, label map, and expected upper-level links. + size_t required_host = estimate_hnsw_host_memory(n_rows, + dim, + graph_degree_int, + cagra_index.metric(), + params.hierarchy, + params.ef_construction); // Honor an explicit host-memory limit from ACE params (if configured), mirroring // hnsw::build. This also makes the spill branch deterministically testable. @@ -1462,6 +1299,21 @@ std::unique_ptr> from_cagra( if (ace.max_host_memory_gb > 0) { max_host_memory_gb = ace.max_host_memory_gb; } } size_t available_host = get_available_memory(max_host_memory_gb).first; + if (max_host_memory_gb.has_value()) { + auto graph = cagra_index.graph(); + cudaPointerAttributes attributes; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attributes, graph.data_handle())); + const bool graph_uses_host_memory = + attributes.type == cudaMemoryTypeUnregistered || attributes.hostPointer != nullptr; + if (graph_uses_host_memory) { + const size_t configured_host = + static_cast(max_host_memory_gb.value() * static_cast(uint64_t{1} << 30)); + const size_t graph_bytes = graph.size() * sizeof(uint32_t); + const size_t remaining_configured_host = + graph_bytes < configured_host ? configured_host - graph_bytes : 0; + available_host = std::min(available_host, remaining_configured_host); + } + } RAFT_LOG_INFO( "hnsw::from_cagra - in-memory HNSW requires ~%4.1f GB host mem, available %4.1f GB", @@ -1670,8 +1522,7 @@ void deserialize(raft::resources const& res, * This function builds an HNSW index * 1. Converting HNSW parameters to CAGRA parameters * 2. Inspect memory requirements (fall back to ACE algorithm if memory constrained) - * 3. Building a CAGRA index with the chosen algorithm in (2) - * 4. Converting the CAGRA index to HNSW format (in-memory or disk-backed) + * 3. Building with ACE (direct partitioned HNSW) or in-memory CAGRA, then converting as needed */ template std::unique_ptr> build(raft::resources const& res, @@ -1679,6 +1530,8 @@ std::unique_ptr> build(raft::resources const& res, raft::host_matrix_view dataset) { common::nvtx::range fun_scope("hnsw::build"); + RAFT_EXPECTS(params.M <= static_cast(std::numeric_limits::max() / 3), + "HNSW M is too large for CAGRA graph-degree parameters"); cuvs::neighbors::cagra::index_params cagra_params = cagra::index_params::from_hnsw_params(dataset.extents(), @@ -1686,12 +1539,18 @@ std::unique_ptr> build(raft::resources const& res, params.ef_construction, cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, params.metric); - cagra_params.metric = params.metric; - - // If the user explicitly configured ACE, honor it. Otherwise (default params) apply a - // heuristic that falls back to ACE only when an in-memory CAGRA build would not fit in - // the available host/device memory. - bool use_ace = std::holds_alternative(params.graph_build_params); + cagra_params.metric = params.metric; + cagra_params.attach_dataset_on_build = false; + + const bool cagra_ace_explicitly_selected = + std::holds_alternative(params.graph_build_params); + if (cagra_ace_explicitly_selected) { + const auto& explicit_ace_params = + std::get(params.graph_build_params); + RAFT_EXPECTS(explicit_ace_params.npartitions <= static_cast(dataset.extent(0)), + "ACE: number of partitions cannot exceed dataset size"); + } + bool cagra_ace_selected_for_memory = false; if (std::holds_alternative(params.graph_build_params)) { auto [required_host, required_dev] = cuvs::neighbors::cagra::helpers::cagra_build_mem_usage( @@ -1707,41 +1566,96 @@ std::unique_ptr> build(raft::resources const& res, if (required_host < available_host && required_dev < available_dev) { RAFT_LOG_INFO("We have sufficient memory to proceed with in memory build"); } else { - use_ace = true; - RAFT_LOG_INFO( - "Not enough host or device memory. Falling back to ACE build with optional disk spilling"); + cagra_ace_selected_for_memory = true; + RAFT_LOG_INFO("Not enough host or device memory. Falling back to ACE partitioned HNSW build"); } } - if (use_ace) { + const bool cagra_ace_selected = cagra_ace_explicitly_selected || cagra_ace_selected_for_memory; + // Partitioned ACE needs more rows than the CAGRA intermediate degree. Smaller ACE requests keep + // the in-memory CAGRA conversion. + const bool ace_partitioned_build_possible = + dataset.extent(0) > static_cast(cagra_params.intermediate_graph_degree); + + if (cagra_ace_selected && ace_partitioned_build_possible) { + RAFT_EXPECTS(params.hierarchy != HnswHierarchy::CPU, + "ACE HNSW construction does not support a CPU hierarchy"); auto ace_params = std::holds_alternative(params.graph_build_params) ? std::get(params.graph_build_params) : graph_build_params::ace_params{}; - // Configure ACE parameters for CAGRA - cuvs::neighbors::cagra::graph_build_params::ace_params cagra_ace_params; - cagra_ace_params.npartitions = ace_params.npartitions; - cagra_ace_params.ef_construction = params.ef_construction; - cagra_ace_params.build_dir = ace_params.build_dir; - cagra_ace_params.use_disk = ace_params.use_disk; - cagra_ace_params.max_host_memory_gb = ace_params.max_host_memory_gb; - cagra_ace_params.max_gpu_memory_gb = ace_params.max_gpu_memory_gb; - cagra_params.graph_build_params = cagra_ace_params; + auto [available_host, available_device] = get_available_memory(); + if (ace_params.max_host_memory_gb > 0) { + available_host = std::min(available_host, + static_cast(ace_params.max_host_memory_gb * + static_cast(uint64_t{1} << 30))); + } + if (ace_params.max_gpu_memory_gb > 0) { + available_device = std::min( + available_device, + static_cast(ace_params.max_gpu_memory_gb * static_cast(uint64_t{1} << 30))); + } + cagra::detail::ace_external_plan_input external_input; + external_input.rows = dataset.extent(0); + external_input.dim = dataset.extent(1); + external_input.element_size = sizeof(T); + external_input.M = params.M; + external_input.intermediate_degree = cagra_params.intermediate_graph_degree; + external_input.graph_degree = cagra_params.graph_degree; + external_input.requested_partitions = ace_params.npartitions; + external_input.available_host_bytes = available_host; + external_input.available_device_bytes = available_device; + uint64_t initial_partitions = + std::min(cagra::detail::external_maximum_partitions( + dataset.extent(0), cagra_params.intermediate_graph_degree), + std::max(2, ace_params.npartitions)); + uint64_t initial_max_occurrences = cagra::detail::external_checked_mul( + 6, + cagra::detail::external_div_rounding_up(dataset.extent(0), initial_partitions), + "initial CAGRA-ACE partition occurrence count"); + auto [optimize_host, optimize_device, optimize_host_fixed, optimize_device_fixed] = + cagra::helpers::optimize_workspace_size(initial_max_occurrences, + cagra_params.graph_degree, + cagra_params.intermediate_graph_degree, + sizeof(uint32_t), + cagra_params.guarantee_connectivity); + external_input.optimize_host_fixed = optimize_host_fixed; + external_input.optimize_device_fixed = optimize_device_fixed; + external_input.optimize_host_per_row = cagra::detail::external_div_rounding_up( + optimize_host - optimize_host_fixed, initial_max_occurrences); + external_input.optimize_device_per_row = cagra::detail::external_div_rounding_up( + optimize_device - optimize_device_fixed, initial_max_occurrences); + external_input.force_disk = true; + external_input.hierarchy = params.hierarchy == HnswHierarchy::GPU; + auto external_plan = cagra::detail::make_ace_external_plan(external_input); + RAFT_LOG_INFO( + "hnsw::build - using ACE partitioned HNSW build with %zu partitions, planned host/device " + "peaks %.3f/%.3f GiB", + static_cast(external_plan.partitions), + external_plan.host_peak_bytes / static_cast(uint64_t{1} << 30), + external_plan.device_peak_bytes / static_cast(uint64_t{1} << 30)); + return external::build_external(res, + params, + dataset, + external_plan, + ace_params, + cagra_params.graph_degree, + cagra_params.intermediate_graph_degree, + params.ef_construction); } // Public HNSW API uses host_matrix_view; CAGRA build expects a padded dataset view. // Host build stores only the graph; vectors are passed separately to from_cagra below. cuvs::neighbors::host_padded_dataset_view host_padded_view( dataset, static_cast(dataset.extent(1))); - auto ace_host_index = cuvs::neighbors::cagra::build(res, cagra_params, host_padded_view); + auto cagra_index = cuvs::neighbors::cagra::build(res, cagra_params, host_padded_view); RAFT_LOG_INFO("hnsw::build - Converting CAGRA index to HNSW format"); - return from_cagra( res, params, - ace_host_index, - ace_host_index.dataset_fd().has_value() ? std::nullopt : std::make_optional(dataset)); + cagra_index, + cagra_index.dataset_fd().has_value() ? std::nullopt : std::make_optional(dataset)); } } // namespace cuvs::neighbors::hnsw::detail diff --git a/cpp/src/neighbors/detail/hnsw/external_build.cuh b/cpp/src/neighbors/detail/hnsw/external_build.cuh new file mode 100644 index 0000000000..422fc94e0f --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/external_build.cuh @@ -0,0 +1,1259 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "../cagra/ace_external_plan.hpp" +#include "external_format.hpp" +#include "external_translate.hpp" +#include "external_workspace.hpp" +#include "index_impl.hpp" +#include "serialize_layout.hpp" + +#include +#include +#include +#include +#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::external { + +struct external_io_counters { + uint64_t sample_read_bytes = 0; + uint64_t source_read_bytes = 0; + uint64_t stage_write_bytes = 0; + uint64_t stage_read_bytes = 0; + uint64_t sidecar_write_bytes = 0; + uint64_t sidecar_read_bytes = 0; + uint64_t output_write_bytes = 0; + uint64_t stage_write_requests = 0; + uint64_t stage_read_requests = 0; + uint64_t sidecar_write_requests = 0; + uint64_t sidecar_read_requests = 0; + uint64_t output_write_requests = 0; +}; + +inline uint64_t elapsed_milliseconds(std::chrono::steady_clock::time_point start, + std::chrono::steady_clock::time_point end) +{ + return static_cast( + std::chrono::duration_cast(end - start).count()); +} + +template +void copy_rows_as_float(const T* source, float* destination, uint64_t rows, uint64_t dimension) +{ + const uint64_t elements = checked_file_mul(rows, dimension, "float conversion elements"); + RAFT_EXPECTS(elements <= static_cast(std::numeric_limits::max()), + "Float conversion exceeds supported element count"); + if constexpr (std::is_same_v) { + std::memcpy(destination, source, static_cast(elements) * sizeof(float)); + } else { +#pragma omp parallel for + for (int64_t index = 0; index < static_cast(elements); ++index) { + destination[index] = static_cast(source[index]); + } + } +} + +template +uint64_t external_parameter_fingerprint(uint64_t rows, + uint64_t dim, + const index_params& params, + const cagra::detail::ace_external_plan& plan, + size_t graph_degree, + size_t intermediate_graph_degree, + size_t ace_ef_construction) +{ + std::array fields{rows, + dim, + sizeof(T), + params.M, + static_cast(params.ef_construction), + ace_ef_construction, + static_cast(params.metric), + static_cast(params.hierarchy), + plan.partitions, + graph_degree, + intermediate_graph_degree, + stage_schema_version}; + return fnv1a64(fields.data(), sizeof(fields)); +} + +template +class stage_buffer_pool { + public: + struct file_state { + std::filesystem::path path; + stage_header header; + uint64_t appended_records = 0; + uint64_t flushed_records = 0; + uint64_t allocated_bytes = stage_data_offset; + std::vector buffer; + size_t used = 0; + uint64_t access = 0; + }; + + stage_buffer_pool(external_workspace& workspace, + uint32_t dimension, + uint32_t partitions, + uint64_t total_buffer_bytes, + uint64_t preferred_buffer_bytes, + uint64_t fingerprint) + : workspace_(workspace), + dimension_(dimension), + partitions_(partitions), + record_size_(checked_file_add(2 * sizeof(uint32_t), + checked_file_mul(dimension, sizeof(T), "stage vector row"), + "stage record")), + growth_extent_bytes_( + std::clamp(checked_file_mul(record_size_, 1024, "stage growth extent"), + uint64_t{64} << 10, + uint64_t{8} << 20)), + total_buffer_bytes_(std::max(total_buffer_bytes, record_size_)), + preferred_buffer_bytes_( + std::max(record_size_, std::min(preferred_buffer_bytes, total_buffer_bytes_))) + { + if (partitions_ == 0) { throw std::invalid_argument("external HNSW requires partitions"); } + // Prefer one buffer per stage file. Fewer LRU buffers make random assignment nearly one write + // per record. + const uint64_t file_count = checked_file_mul(partitions_, 2, "external stage file count"); + const uint64_t per_file_budget = + std::max(record_size_, total_buffer_bytes_ / file_count); + preferred_buffer_bytes_ = std::min(preferred_buffer_bytes_, per_file_budget); + files_.reserve(static_cast(partitions) * 2); + for (uint32_t partition = 0; partition < partitions; ++partition) { + add_file(partition, false, fingerprint); + add_file(partition, true, fingerprint); + } + } + + ~stage_buffer_pool() noexcept + { + try { + flush_all(); + } catch (...) { + } + } + + void append_core(uint32_t partition, uint32_t original_label, const T* vector) + { + append(file_index(partition, false), original_label, 0, vector); + } + + void append_spill(uint32_t partition, + uint32_t owner_partition, + uint32_t owner_ordinal, + const T* vector) + { + append(file_index(partition, true), owner_partition, owner_ordinal, vector); + } + + void finalize() + { + flush_all(); + for (auto& state : files_) { + cuvs::util::file_descriptor fd(state.path.string(), O_RDWR); + state.header.committed_records = state.appended_records; + uint64_t exact_size = expected_file_size(state.header); + if (exact_size > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("external HNSW stage exceeds off_t"); + } + if (::ftruncate(fd.get(), static_cast(exact_size)) != 0) { + throw std::runtime_error("failed to trim external HNSW stage: " + + std::string{strerror(errno)}); + } + write_stage_header(fd.get(), state.header); + if (::fsync(fd.get()) != 0) { + throw std::runtime_error("failed to commit external HNSW stage: " + + std::string{strerror(errno)}); + } + validate_stage_file_size(fd.get(), state.header); + std::vector{}.swap(state.buffer); + } + resident_buffer_bytes_ = 0; + } + + [[nodiscard]] const file_state& core(uint32_t partition) const + { + return files_.at(file_index(partition, false)); + } + [[nodiscard]] const file_state& spill(uint32_t partition) const + { + return files_.at(file_index(partition, true)); + } + [[nodiscard]] uint64_t record_size() const noexcept { return record_size_; } + [[nodiscard]] uint64_t bytes_written() const noexcept { return bytes_written_; } + [[nodiscard]] uint64_t write_requests() const noexcept { return write_requests_; } + + private: + [[nodiscard]] size_t file_index(uint32_t partition, bool spill) const + { + if (partition >= partitions_) { throw std::out_of_range("external HNSW partition"); } + return static_cast(partition) * 2 + static_cast(spill); + } + + void add_file(uint32_t partition, bool spill, uint64_t fingerprint) + { + std::ostringstream filename; + filename << (spill ? "spill." : "core.") << partition << ".stage"; + auto path = workspace_.private_path(filename.str()); + workspace_.create_stage_file(filename.str(), + spill ? stage_kind::spill : stage_kind::core, + dimension_, + partition, + 0, + record_size_); + file_state state; + state.path = std::move(path); + state.header = make_stage_header(spill ? stage_kind::spill : stage_kind::core, + dimension_, + partition, + 0, + record_size_, + fingerprint); + files_.push_back(std::move(state)); + } + + void ensure_buffer(size_t index) + { + auto& state = files_[index]; + if (!state.buffer.empty()) { return; } + while (resident_buffer_bytes_ + preferred_buffer_bytes_ > total_buffer_bytes_) { + size_t victim = files_.size(); + for (size_t candidate = 0; candidate < files_.size(); ++candidate) { + if (!files_[candidate].buffer.empty() && candidate != index && + (victim == files_.size() || files_[candidate].access < files_[victim].access)) { + victim = candidate; + } + } + if (victim == files_.size()) { break; } + flush(victim); + resident_buffer_bytes_ -= files_[victim].buffer.size(); + std::vector{}.swap(files_[victim].buffer); + } + state.buffer.resize(static_cast(preferred_buffer_bytes_)); + resident_buffer_bytes_ += state.buffer.size(); + } + + void append(size_t index, uint32_t first, uint32_t second, const T* vector) + { + ensure_buffer(index); + auto& state = files_[index]; + if (state.buffer.size() - state.used < record_size_) { flush(index); } + std::byte* destination = state.buffer.data() + state.used; + std::memcpy(destination, &first, sizeof(first)); + std::memcpy(destination + sizeof(first), &second, sizeof(second)); + std::memcpy( + destination + 2 * sizeof(uint32_t), vector, static_cast(dimension_) * sizeof(T)); + state.used += static_cast(record_size_); + ++state.appended_records; + state.access = ++clock_; + } + + void reserve_extent(file_state& state, int fd, uint64_t required_size) + { + if (required_size <= state.allocated_bytes) { return; } + uint64_t rounded = + checked_file_add(required_size, growth_extent_bytes_ - 1, "rounded stage growth extent"); + uint64_t new_size = + checked_file_mul(rounded / growth_extent_bytes_, growth_extent_bytes_, "stage growth extent"); + if (new_size > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("external HNSW stage allocation exceeds off_t"); + } + int status = ::posix_fallocate(fd, + static_cast(state.allocated_bytes), + static_cast(new_size - state.allocated_bytes)); + if (status != 0) { + throw std::runtime_error("failed to grow external HNSW stage: " + + std::string{strerror(status)}); + } + state.allocated_bytes = new_size; + } + + void flush(size_t index) + { + auto& state = files_[index]; + if (state.used == 0) { return; } + cuvs::util::file_descriptor fd(state.path.string(), O_RDWR); + uint64_t offset = + checked_file_add(stage_data_offset, + checked_file_mul(state.flushed_records, record_size_, "stage append offset"), + "stage append offset"); + uint64_t end = checked_file_add(offset, state.used, "stage append end"); + reserve_extent(state, fd.get(), end); + cuvs::util::write_large_file(fd, state.buffer.data(), state.used, offset); + state.flushed_records += state.used / record_size_; + bytes_written_ += state.used; + ++write_requests_; + state.used = 0; + } + + void flush_all() + { + for (size_t index = 0; index < files_.size(); ++index) { + flush(index); + } + } + + external_workspace& workspace_; + uint32_t dimension_; + uint32_t partitions_; + uint64_t record_size_; + uint64_t growth_extent_bytes_; + uint64_t total_buffer_bytes_; + uint64_t preferred_buffer_bytes_; + std::vector files_; + uint64_t resident_buffer_bytes_ = 0; + uint64_t bytes_written_ = 0; + uint64_t write_requests_ = 0; + uint64_t clock_ = 0; +}; + +class buffered_stage_reader { + public: + buffered_stage_reader(const cuvs::util::file_descriptor& fd, + stage_header header, + size_t target_buffer_bytes) + : fd_(fd), + header_(header), + records_per_buffer_( + std::max(1, target_buffer_bytes / std::max(1, header.record_size))), + buffer_(static_cast( + checked_file_mul(records_per_buffer_, header.record_size, "stage read buffer"))) + { + const auto path = fd_.get_path(); + if (!path.empty()) { file_handle_ = std::make_unique(path, "r"); } + } + + bool next(const std::byte*& record) + { + if (current_record_ == header_.committed_records) { return false; } + if (buffer_index_ == buffered_records_) { refill(); } + record = buffer_.data() + static_cast(buffer_index_ * header_.record_size); + ++buffer_index_; + ++current_record_; + return true; + } + + [[nodiscard]] uint64_t bytes_read() const noexcept { return bytes_read_; } + [[nodiscard]] uint64_t read_requests() const noexcept { return read_requests_; } + + private: + void refill() + { + uint64_t remaining = header_.committed_records - current_record_; + buffered_records_ = std::min(records_per_buffer_, remaining); + size_t bytes = + static_cast(checked_file_mul(buffered_records_, header_.record_size, "stage read")); + uint64_t offset = + checked_file_add(header_.data_offset, + checked_file_mul(current_record_, header_.record_size, "stage read offset"), + "stage read offset"); + if (file_handle_) { + const size_t bytes_read = file_handle_->pread(buffer_.data(), bytes, offset).get(); + RAFT_EXPECTS(bytes_read == bytes, + "Incomplete stage read: expected %zu bytes, read %zu", + bytes, + bytes_read); + } else { + cuvs::util::read_large_file(fd_, buffer_.data(), bytes, offset); + } + bytes_read_ += bytes; + ++read_requests_; + buffer_index_ = 0; + } + + const cuvs::util::file_descriptor& fd_; + stage_header header_; + uint64_t records_per_buffer_; + std::vector buffer_; + std::unique_ptr file_handle_; + uint64_t current_record_ = 0; + uint64_t buffered_records_ = 0; + uint64_t buffer_index_ = 0; + uint64_t bytes_read_ = 0; + uint64_t read_requests_ = 0; +}; + +template +raft::device_matrix train_external_centroids( + raft::resources const& res, + raft::host_matrix_view dataset, + const cagra::detail::ace_external_plan& plan, + external_io_counters& io) +{ + const uint64_t rows = dataset.extent(0); + const uint64_t dim = dataset.extent(1); + const uint64_t sample_rows = plan.centroid_sample_rows; + auto sample = raft::make_host_matrix(sample_rows, dim); + + uint64_t copied = 0; + for (const auto& range : cagra::detail::make_external_sample_ranges(rows, sample_rows)) { + copy_rows_as_float(dataset.data_handle() + range.start * dim, + sample.data_handle() + copied * dim, + range.count, + dim); + copied += range.count; + } + RAFT_EXPECTS(copied == sample_rows, "External HNSW centroid sample size mismatch"); + io.sample_read_bytes = checked_file_mul( + sample_rows, checked_file_mul(dim, sizeof(T), "sample vector"), "sample bytes"); + + auto sample_device = raft::make_device_matrix(res, sample_rows, dim); + raft::copy(res, sample_device.view(), sample.view()); + auto centroids = + raft::make_device_matrix(res, plan.partitions, dataset.extent(1)); + cuvs::cluster::kmeans::balanced_params kmeans_params; + cuvs::cluster::kmeans::fit(res, kmeans_params, sample_device.view(), centroids.view()); + return centroids; +} + +template +void assign_and_stage(raft::resources const& res, + raft::host_matrix_view dataset, + const cagra::detail::ace_external_plan& plan, + raft::device_matrix_view centroids, + stage_buffer_pool& stages, + external_io_counters& io) +{ + const uint64_t rows = dataset.extent(0); + const uint64_t dim = dataset.extent(1); + const uint64_t partitions = plan.partitions; + const uint64_t chunk_rows = plan.assignment_chunk_rows; + + auto host_input = raft::make_host_matrix(chunk_rows, dim); + auto device_input = raft::make_device_matrix(res, chunk_rows, dim); + auto device_distances = raft::make_device_matrix(res, chunk_rows, partitions); + auto device_top_distances = raft::make_device_matrix(res, chunk_rows, 2); + auto device_top_labels = raft::make_device_matrix(res, chunk_rows, 2); + auto host_top_labels = raft::make_host_matrix(chunk_rows, 2); + std::vector core_counts(partitions, 0); + + for (uint64_t base = 0; base < rows; base += chunk_rows) { + uint64_t count = std::min(chunk_rows, rows - base); + auto host_view = + raft::make_host_matrix_view(host_input.data_handle(), count, dim); + copy_rows_as_float(dataset.data_handle() + base * dim, host_view.data_handle(), count, dim); + auto device_view = + raft::make_device_matrix_view(device_input.data_handle(), count, dim); + auto distance_view = raft::make_device_matrix_view( + device_distances.data_handle(), count, partitions); + auto top_distance_view = + raft::make_device_matrix_view(device_top_distances.data_handle(), count, 2); + auto top_label_view = + raft::make_device_matrix_view(device_top_labels.data_handle(), count, 2); + auto host_label_view = + raft::make_host_matrix_view(host_top_labels.data_handle(), count, 2); + raft::copy(res, device_view, host_view); + cuvs::distance::pairwise_distance(res, + raft::make_const_mdspan(device_view), + centroids, + distance_view, + cuvs::distance::DistanceType::L2Expanded); + cuvs::selection::select_k(res, + raft::make_const_mdspan(distance_view), + std::nullopt, + top_distance_view, + top_label_view, + true, + true); + raft::copy(res, host_label_view, top_label_view); + raft::resource::sync_stream(res); + + for (uint64_t local = 0; local < count; ++local) { + uint32_t core_partition = host_label_view(local, 0); + uint32_t spill_partition = host_label_view(local, 1); + RAFT_EXPECTS(core_partition < partitions && spill_partition < partitions && + core_partition != spill_partition, + "Invalid external HNSW top-two partition assignment"); + uint64_t ordinal = core_counts[core_partition]++; + RAFT_EXPECTS(ordinal <= std::numeric_limits::max(), + "External HNSW core ordinal exceeds uint32_t"); + auto* vector = dataset.data_handle() + (base + local) * dim; + stages.append_core(core_partition, static_cast(base + local), vector); + stages.append_spill(spill_partition, core_partition, static_cast(ordinal), vector); + } + } + io.source_read_bytes = + checked_file_mul(rows, checked_file_mul(dim, sizeof(T), "source row"), "source scan bytes"); + io.stage_write_bytes = stages.bytes_written(); + io.stage_write_requests = stages.write_requests(); +} + +template +struct resident_partition { + raft::host_matrix vectors; + std::vector local_to_global; + std::vector core_labels; + uint64_t core_rows = 0; +}; + +inline uint32_t external_owner_id(const std::vector& prefixes, + uint32_t owner_partition, + uint32_t owner_ordinal) +{ + uint64_t owner_next = static_cast(owner_partition) + 1; + RAFT_EXPECTS(owner_next < prefixes.size(), "Invalid external HNSW spill owner partition"); + uint64_t owner_begin = prefixes[owner_partition]; + uint64_t owner_end = prefixes[owner_next]; + uint64_t global = checked_file_add(owner_begin, owner_ordinal, "external HNSW owner ID"); + RAFT_EXPECTS(global < owner_end, "Invalid external HNSW spill owner ordinal"); + return static_cast(global); +} + +template +resident_partition read_resident_partition( + const typename stage_buffer_pool::file_state& core_state, + const typename stage_buffer_pool::file_state& spill_state, + uint32_t partition, + const std::vector& prefixes, + uint64_t expected_record_size, + uint64_t fingerprint, + uint64_t reader_buffer_bytes, + external_io_counters& io) +{ + cuvs::util::file_descriptor core_fd(core_state.path.string(), O_RDONLY); + cuvs::util::file_descriptor spill_fd(spill_state.path.string(), O_RDONLY); + auto core_header = read_stage_header(core_fd.get()); + auto spill_header = read_stage_header(spill_fd.get()); + validate_stage_header(core_header, + stage_kind::core, + element_kind_for(), + sizeof(T), + core_header.dimension, + partition, + 0, + expected_record_size, + fingerprint); + validate_stage_header(spill_header, + stage_kind::spill, + element_kind_for(), + sizeof(T), + spill_header.dimension, + partition, + 0, + expected_record_size, + fingerprint); + RAFT_EXPECTS(core_header.dimension == spill_header.dimension, + "External HNSW core/spill dimension mismatch"); + validate_stage_file_size(core_fd.get(), core_header); + validate_stage_file_size(spill_fd.get(), spill_header); + + uint64_t total_rows = checked_file_add( + core_header.committed_records, spill_header.committed_records, "resident partition rows"); + resident_partition resident{ + raft::make_host_matrix(total_rows, core_header.dimension), + std::vector(total_rows), + std::vector(core_header.committed_records), + core_header.committed_records}; + + auto read_file = [&](const cuvs::util::file_descriptor& fd, + const stage_header& header, + bool spill, + uint64_t destination_base) { + buffered_stage_reader reader( + fd, header, static_cast(std::max(1, reader_buffer_bytes))); + const std::byte* record = nullptr; + uint64_t row = 0; + while (reader.next(record)) { + uint32_t first; + uint32_t second; + std::memcpy(&first, record, sizeof(first)); + std::memcpy(&second, record + sizeof(first), sizeof(second)); + std::memcpy(&resident.vectors(destination_base + row, 0), + record + 2 * sizeof(uint32_t), + static_cast(header.dimension) * sizeof(T)); + if (spill) { + resident.local_to_global[destination_base + row] = + external_owner_id(prefixes, first, second); + } else { + uint64_t global = static_cast(prefixes[partition]) + row; + RAFT_EXPECTS(global < prefixes.back(), "Invalid external HNSW core prefix"); + resident.local_to_global[destination_base + row] = static_cast(global); + resident.core_labels[row] = first; + } + ++row; + } + RAFT_EXPECTS(row == header.committed_records, "External HNSW stage record count mismatch"); + io.stage_read_bytes += reader.bytes_read(); + io.stage_read_requests += reader.read_requests(); + }; + read_file(core_fd, core_header, false, 0); + read_file(spill_fd, spill_header, true, resident.core_rows); + return resident; +} + +template +class level_file { + public: + level_file(external_workspace& workspace, + std::string filename, + stage_kind kind, + uint32_t dimension, + uint32_t level, + uint64_t record_size, + uint64_t fingerprint, + size_t buffer_size) + : fd_(workspace.create_stage_file(filename, kind, dimension, 0, level, record_size)), + header_(make_stage_header(kind, dimension, 0, level, record_size, fingerprint)), + writer_(fd_, stage_data_offset, buffer_size) + { + } + + void append(const void* record) + { + writer_.write(record, static_cast(header_.record_size)); + ++count_; + } + + void finalize() + { + writer_.flush(); + header_.committed_records = count_; + uint64_t exact = expected_file_size(header_); + if (::ftruncate(fd_.get(), static_cast(exact)) != 0) { + throw std::runtime_error("failed to finalize external HNSW level file"); + } + write_stage_header(fd_.get(), header_); + if (::fsync(fd_.get()) != 0) { + throw std::runtime_error("failed to sync external HNSW level file"); + } + validate_stage_file_size(fd_.get(), header_); + } + + [[nodiscard]] const cuvs::util::file_descriptor& descriptor() const noexcept { return fd_; } + [[nodiscard]] const stage_header& header() const noexcept { return header_; } + [[nodiscard]] uint64_t count() const noexcept { return count_; } + [[nodiscard]] uint64_t request_count() const noexcept { return writer_.request_count(); } + + private: + cuvs::util::file_descriptor fd_; + stage_header header_; + sequential_file_writer writer_; + uint64_t count_ = 0; +}; + +template +void write_partitioned_level(raft::resources const& res, + const resident_partition& resident, + const std::vector& hierarchy_levels, + uint32_t level, + const hnsw_serialize_layout& layout, + cuvs::distance::DistanceType metric, + level_file& sidecar) +{ + std::vector promoted_rows; + promoted_rows.reserve(resident.local_to_global.size() / std::max(2, layout.M)); + for (uint64_t row = 0; row < resident.local_to_global.size(); ++row) { + if (hierarchy_levels[row] >= level) { promoted_rows.push_back(row); } + } + + const uint64_t count = promoted_rows.size(); + auto promoted_vectors = + raft::make_host_matrix(count, static_cast(layout.dimension)); + std::vector promoted_ids(count); + std::vector core_promoted_position(resident.core_rows, -1); + for (uint64_t promoted = 0; promoted < count; ++promoted) { + uint64_t source = promoted_rows[promoted]; + std::copy(&resident.vectors(source, 0), + &resident.vectors(source, 0) + layout.dimension, + &promoted_vectors(promoted, 0)); + promoted_ids[promoted] = resident.local_to_global[source]; + if (source < resident.core_rows) { core_promoted_position[source] = promoted; } + } + + uint64_t neighbor_count = count > 1 ? std::min(layout.M, count - 1) : 0; + auto neighbors = raft::make_host_matrix(count, neighbor_count); + if (count > 1) { + all_neighbors_graph( + res, raft::make_const_mdspan(promoted_vectors.view()), neighbors.view(), metric); + } + + std::vector record(1 + layout.M, std::numeric_limits::max()); + for (uint64_t core = 0; core < resident.core_rows; ++core) { + int64_t promoted = core_promoted_position[core]; + if (promoted < 0) { continue; } + record[0] = resident.local_to_global[core]; + std::fill(record.begin() + 1, record.end(), std::numeric_limits::max()); + for (uint64_t neighbor = 0; neighbor < neighbor_count; ++neighbor) { + uint32_t local = neighbors(promoted, neighbor); + RAFT_EXPECTS(local < promoted_ids.size(), "Invalid external HNSW upper local ID"); + record[neighbor + 1] = promoted_ids[local]; + } + sidecar.append(record.data()); + } +} + +template +void append_global_level_vectors(const resident_partition& resident, + const std::vector& hierarchy_levels, + uint32_t level, + const hnsw_serialize_layout& layout, + level_file& vectors) +{ + std::vector record(sizeof(uint32_t) + layout.dimension * sizeof(T)); + for (uint64_t core = 0; core < resident.core_rows; ++core) { + uint32_t id = resident.local_to_global[core]; + if (hierarchy_levels[core] < level) { continue; } + std::memcpy(record.data(), &id, sizeof(id)); + std::memcpy( + record.data() + sizeof(id), &resident.vectors(core, 0), layout.dimension * sizeof(T)); + vectors.append(record.data()); + } +} + +template +void build_global_level(raft::resources const& res, + level_file& vectors, + level_file& sidecar, + const hnsw_serialize_layout& layout, + cuvs::distance::DistanceType metric, + uint64_t reader_buffer_bytes, + external_io_counters& io) +{ + vectors.finalize(); + const auto& header = vectors.header(); + buffered_stage_reader reader(vectors.descriptor(), header, reader_buffer_bytes); + auto dataset = raft::make_host_matrix(header.committed_records, layout.dimension); + std::vector ids(header.committed_records); + const std::byte* record = nullptr; + uint64_t row = 0; + while (reader.next(record)) { + std::memcpy(&ids[row], record, sizeof(uint32_t)); + std::memcpy(&dataset(row, 0), record + sizeof(uint32_t), layout.dimension * sizeof(T)); + ++row; + } + io.sidecar_read_bytes += reader.bytes_read(); + io.sidecar_read_requests += reader.read_requests(); + + uint64_t neighbor_count = + row > 1 ? std::min(layout.M, static_cast(row - 1)) : 0; + auto neighbors = raft::make_host_matrix(row, neighbor_count); + if (row > 1) { + all_neighbors_graph(res, raft::make_const_mdspan(dataset.view()), neighbors.view(), metric); + } + std::vector output(1 + layout.M, std::numeric_limits::max()); + for (uint64_t index = 0; index < row; ++index) { + output[0] = ids[index]; + std::fill(output.begin() + 1, output.end(), std::numeric_limits::max()); + for (uint64_t neighbor = 0; neighbor < neighbor_count; ++neighbor) { + uint32_t local = neighbors(index, neighbor); + RAFT_EXPECTS(local < ids.size(), "Invalid external HNSW global upper local ID"); + output[neighbor + 1] = ids[local]; + } + sidecar.append(output.data()); + } +} + +template +std::unique_ptr> build_external( + raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset, + const cagra::detail::ace_external_plan& plan, + const graph_build_params::ace_params& ace_params, + size_t graph_degree, + size_t intermediate_graph_degree, + size_t ace_ef_construction) +{ + RAFT_EXPECTS(plan.use_disk, "External HNSW builder requires a disk plan"); + RAFT_EXPECTS(params.hierarchy != HnswHierarchy::CPU, + "Disk HNSW construction does not support a CPU hierarchy"); + RAFT_EXPECTS(graph_degree == 2 * params.M, + "External HNSW graph degree must equal the base-layer capacity 2*M"); + RAFT_EXPECTS(intermediate_graph_degree >= graph_degree, + "External HNSW intermediate graph degree must be at least graph degree"); + RAFT_EXPECTS(ace_ef_construction > 0, "External HNSW CAGRA-ACE ef_construction must be positive"); + const uint64_t rows = dataset.extent(0); + const uint64_t dim = dataset.extent(1); + const uint64_t fingerprint = external_parameter_fingerprint( + rows, dim, params, plan, graph_degree, intermediate_graph_degree, ace_ef_construction); + [[maybe_unused]] const int num_threads = + params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; + external_workspace workspace(ace_params.build_dir, fingerprint); + external_io_counters io; + auto build_start = std::chrono::steady_clock::now(); + auto preflight_start = build_start; + uint64_t preflight_ms = 0; + uint64_t centroid_training_ms = 0; + uint64_t assignment_stage_ms = 0; + uint64_t partition_build_ms = 0; + uint64_t partition_load_ms = 0; + uint64_t cagra_build_ms = 0; + uint64_t base_serialization_ms = 0; + uint64_t hierarchy_ms = 0; + uint64_t publication_ms = 0; + + try { + auto hierarchy = params.hierarchy == HnswHierarchy::GPU ? summarize_hierarchy(rows, params.M) + : hierarchy_summary{}; + auto hnsw_index = + std::make_unique>(static_cast(dim), params.metric, params.hierarchy); + auto layout = make_hnsw_serialize_layout::type>( + hnsw_index->get_space(), + rows, + dim, + graph_degree, + params.ef_construction, + params.hierarchy == HnswHierarchy::GPU, + hierarchy); + uint64_t exact_output = layout.exact_file_size(hierarchy.total_active_occurrences); + auto output_fd = workspace.create_partial(exact_output); + auto free_space = std::filesystem::space(workspace.invocation_dir()); + uint64_t stage_record_size = checked_file_add( + 2 * sizeof(uint32_t), checked_file_mul(dim, sizeof(T), "stage vector"), "stage record"); + uint64_t stage_growth_extent = + std::clamp(checked_file_mul(stage_record_size, 1024, "stage growth extent"), + uint64_t{64} << 10, + uint64_t{8} << 20); + uint64_t stage_extent_overhead = checked_file_mul( + checked_file_mul(plan.partitions, 2, "stage file count"), + checked_file_add(stage_data_offset, stage_growth_extent, "stage file overhead"), + "stage extent overhead"); + uint64_t required_sidecar_space = 0; + for (uint32_t level = 1; level <= hierarchy.max_level; ++level) { + const uint64_t active = hierarchy.active_by_level[level - 1]; + const uint64_t link_record = + checked_file_add(sizeof(uint32_t), + checked_file_mul(layout.M, sizeof(uint32_t), "upper links"), + "upper link record"); + required_sidecar_space = checked_file_add( + required_sidecar_space, + checked_file_add(stage_data_offset, + checked_file_mul(active, link_record, "upper link sidecar"), + "upper link sidecar file"), + "upper sidecar high-water space"); + if (active <= plan.global_upper_level_max_rows) { + const uint64_t vector_record = + checked_file_add(sizeof(uint32_t), + checked_file_mul(dim, sizeof(T), "upper vector"), + "upper vector record"); + required_sidecar_space = checked_file_add( + required_sidecar_space, + checked_file_add(stage_data_offset, + checked_file_mul(active, vector_record, "upper vector sidecar"), + "upper vector sidecar file"), + "upper sidecar high-water space"); + } + } + uint64_t required_working_space = checked_file_add( + checked_file_add(plan.bytes.stage_write, stage_extent_overhead, "stage high-water space"), + required_sidecar_space, + "external HNSW working-space high-water mark"); + RAFT_EXPECTS(free_space.available >= required_working_space, + "External HNSW staging and hierarchy sidecars require %zu bytes after final-index " + "preallocation, but only %zu bytes are available", + static_cast(required_working_space), + static_cast(free_space.available)); + preflight_ms = elapsed_milliseconds(preflight_start, std::chrono::steady_clock::now()); + + stage_buffer_pool stages(workspace, + static_cast(dim), + static_cast(plan.partitions), + plan.staging_buffer_bytes, + plan.preferred_buffer_bytes, + fingerprint); + workspace.update_manifest("centroid_training"); + auto centroid_start = std::chrono::steady_clock::now(); + auto assignment_start = centroid_start; + { + auto centroids = train_external_centroids(res, dataset, plan, io); + centroid_training_ms = elapsed_milliseconds(centroid_start, std::chrono::steady_clock::now()); + + workspace.update_manifest("assignment_and_staging"); + assignment_start = std::chrono::steady_clock::now(); + assign_and_stage(res, dataset, plan, raft::make_const_mdspan(centroids.view()), stages, io); + } + stages.finalize(); + io.stage_write_bytes = stages.bytes_written(); + io.stage_write_requests = stages.write_requests(); + assignment_stage_ms = elapsed_milliseconds(assignment_start, std::chrono::steady_clock::now()); + + std::vector prefixes(plan.partitions + 1, 0); + uint64_t total_core = 0; + uint64_t total_spill = 0; + for (uint32_t partition = 0; partition < plan.partitions; ++partition) { + uint64_t core = stages.core(partition).header.committed_records; + uint64_t spill = stages.spill(partition).header.committed_records; + uint64_t actual_occurrences = + checked_file_add(core, spill, "external HNSW actual partition rows"); + uint64_t required_host = checked_file_add( + plan.host_fixed_bytes, + checked_file_mul( + actual_occurrences, plan.host_per_occurrence, "external HNSW actual host partition"), + "external HNSW required host bytes"); + uint64_t required_device = checked_file_add( + plan.device_fixed_bytes, + checked_file_mul( + actual_occurrences, plan.device_per_occurrence, "external HNSW actual device partition"), + "external HNSW required device bytes"); + RAFT_EXPECTS(core + spill <= plan.max_occurrences, + "External HNSW partition %u has %zu rows; planned maximum is %zu. Estimated " + "required/available host bytes are %zu/%zu and device bytes are %zu/%zu. " + "Increase partitions or memory budget.", + partition, + static_cast(actual_occurrences), + static_cast(plan.max_occurrences), + static_cast(required_host), + static_cast(plan.host_budget_bytes), + static_cast(required_device), + static_cast(plan.device_budget_bytes)); + total_core += core; + total_spill += spill; + RAFT_EXPECTS(total_core <= std::numeric_limits::max(), + "External HNSW core prefix exceeds uint32_t"); + prefixes[partition + 1] = static_cast(total_core); + } + RAFT_EXPECTS(total_core == rows && total_spill == rows, + "External HNSW staging counts do not match the dataset"); + workspace.update_manifest("staged", total_core, total_spill); + + sequential_file_writer output(output_fd, 0, static_cast(plan.hnsw_output_buffer_bytes)); + write_hnsw_header(output, layout); + + std::vector>> sidecars(hierarchy.max_level); + std::vector>> global_vectors(hierarchy.max_level); + std::vector global_level(hierarchy.max_level, false); + size_t level_buffer_size = static_cast(std::max( + uint64_t{4} << 10, + std::min( + plan.preferred_buffer_bytes, + plan.staging_buffer_bytes / std::max(1, 2 * hierarchy.max_level)))); + for (uint32_t level = 1; level <= hierarchy.max_level; ++level) { + std::ostringstream sidecar_name; + sidecar_name << "upper." << level << ".links"; + sidecars[level - 1] = + std::make_unique>(workspace, + sidecar_name.str(), + stage_kind::upper_links, + 0, + level, + sizeof(uint32_t) + layout.M * sizeof(uint32_t), + fingerprint, + level_buffer_size); + global_level[level - 1] = + hierarchy.active_by_level[level - 1] <= plan.global_upper_level_max_rows; + if (global_level[level - 1]) { + std::ostringstream vectors_name; + vectors_name << "upper." << level << ".vectors"; + global_vectors[level - 1] = + std::make_unique>(workspace, + vectors_name.str(), + stage_kind::upper_vectors, + static_cast(dim), + level, + sizeof(uint32_t) + dim * sizeof(T), + fingerprint, + level_buffer_size); + } + } + + workspace.update_manifest("partition_build", total_core, total_spill); + auto partition_start = std::chrono::steady_clock::now(); + struct load_result { + resident_partition resident; + external_io_counters io; + uint64_t elapsed_ms; + }; + auto load_partition = [&](uint32_t partition) -> load_result { + auto load_start = std::chrono::steady_clock::now(); + external_io_counters partition_io; + auto resident = read_resident_partition(stages.core(partition), + stages.spill(partition), + partition, + prefixes, + stages.record_size(), + fingerprint, + plan.preferred_buffer_bytes, + partition_io); + return {std::move(resident), + partition_io, + elapsed_milliseconds(load_start, std::chrono::steady_clock::now())}; + }; + std::optional loaded; + if (plan.partitions > 0) { loaded.emplace(load_partition(0)); } + for (uint32_t partition = 0; partition < plan.partitions; ++partition) { + RAFT_EXPECTS(loaded.has_value(), "External HNSW partition prefetch state is empty"); + std::future next; + if (plan.queue_depth > 1 && partition + 1 < plan.partitions) { + next = std::async(std::launch::async, load_partition, partition + 1); + } + + auto current = std::move(*loaded); + loaded.reset(); + partition_load_ms += current.elapsed_ms; + io.stage_read_bytes += current.io.stage_read_bytes; + io.stage_read_requests += current.io.stage_read_requests; + auto& resident = current.resident; + if (resident.core_rows != 0) { + RAFT_EXPECTS(resident.vectors.extent(0) > 1, + "External HNSW partition is too small for graph construction"); + + cagra::index_params partition_params = + cagra::index_params::from_hnsw_params(resident.vectors.view().extents(), + params.M, + ace_ef_construction, + cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, + params.metric); + partition_params.graph_degree = graph_degree; + partition_params.intermediate_graph_degree = intermediate_graph_degree; + RAFT_EXPECTS( + resident.vectors.extent(0) > + static_cast(partition_params.intermediate_graph_degree), + "External HNSW partition %u has %zu rows, but at least %zu are required for M=%zu", + partition, + static_cast(resident.vectors.extent(0)), + partition_params.intermediate_graph_degree + 1, + params.M); + partition_params.attach_dataset_on_build = false; + auto partition_dataset = cuvs::neighbors::make_host_standard_dataset_view( + raft::make_const_mdspan(resident.vectors.view())); + auto cagra_start = std::chrono::steady_clock::now(); + auto partition_index = cagra::build(res, partition_params, partition_dataset); + cagra_build_ms += elapsed_milliseconds(cagra_start, std::chrono::steady_clock::now()); + auto serialization_start = std::chrono::steady_clock::now(); + + std::vector hierarchy_levels; + if (params.hierarchy == HnswHierarchy::GPU) { + hierarchy_levels.resize(resident.local_to_global.size()); +#pragma omp parallel for num_threads(num_threads) + for (int64_t row = 0; row < static_cast(resident.local_to_global.size()); + ++row) { + auto level = + level_for_internal_id(resident.local_to_global[row], hnsw_level_seed, layout.M); + hierarchy_levels[row] = static_cast(level); + } + } + + const uint64_t actual_graph_degree = partition_index.graph().extent(1); + RAFT_EXPECTS(actual_graph_degree == layout.graph_degree, + "External HNSW partition graph degree changed unexpectedly"); + const uint64_t graph_chunk_rows = + std::max(1, + plan.hnsw_output_buffer_bytes / + std::max(1, actual_graph_degree * sizeof(uint32_t))); + auto graph_chunk = + raft::make_host_matrix(graph_chunk_rows, actual_graph_degree); + auto device_graph_chunk = + raft::make_device_matrix(res, graph_chunk_rows, actual_graph_degree); + auto device_mapping = + raft::make_device_vector(res, resident.local_to_global.size()); + raft::copy(res, + device_mapping.view(), + raft::make_host_vector_view( + resident.local_to_global.data(), resident.local_to_global.size())); + for (uint64_t base = 0; base < resident.core_rows; base += graph_chunk_rows) { + uint64_t count = std::min(graph_chunk_rows, resident.core_rows - base); + auto host_view = raft::make_host_matrix_view( + graph_chunk.data_handle(), count, actual_graph_degree); + auto source_view = raft::make_device_matrix_view( + partition_index.graph().data_handle() + base * actual_graph_degree, + count, + actual_graph_degree); + auto translated_view = raft::make_device_matrix_view( + device_graph_chunk.data_handle(), count, actual_graph_degree); + raft::copy(res, translated_view, source_view); + translate_graph_ids(res, + translated_view.data_handle(), + translated_view.size(), + device_mapping.data_handle(), + device_mapping.size()); + raft::copy(res, host_view, raft::make_const_mdspan(translated_view)); + raft::resource::sync_stream(res); + for (uint64_t local = 0; local < count; ++local) { + for (uint64_t edge = 0; edge < actual_graph_degree; ++edge) { + uint32_t neighbor = host_view(local, edge); + RAFT_EXPECTS(neighbor < rows, "External HNSW partition neighbor is out of range"); + } + uint64_t core_row = base + local; + write_hnsw_base_row(output, + layout, + &host_view(local, 0), + &resident.vectors(core_row, 0), + resident.core_labels[core_row]); + } + } + + for (uint32_t level = 1; level <= hierarchy.max_level; ++level) { + if (global_level[level - 1]) { + append_global_level_vectors( + resident, hierarchy_levels, level, layout, *global_vectors[level - 1]); + } else { + write_partitioned_level( + res, resident, hierarchy_levels, level, layout, params.metric, *sidecars[level - 1]); + } + } + base_serialization_ms += + elapsed_milliseconds(serialization_start, std::chrono::steady_clock::now()); + } + + if (partition + 1 < plan.partitions) { + loaded.emplace(plan.queue_depth > 1 ? next.get() : load_partition(partition + 1)); + } + } + + auto base_commit_start = std::chrono::steady_clock::now(); + output.flush(); + RAFT_EXPECTS(output.position() == hnswlib_header_size + layout.base_section_size(), + "External HNSW base section size mismatch"); + if (::fsync(output_fd.get()) != 0) { + throw std::runtime_error("failed to commit external HNSW base section"); + } + base_serialization_ms += + elapsed_milliseconds(base_commit_start, std::chrono::steady_clock::now()); + workspace.update_manifest("base_committed", total_core, total_spill, output.position()); + for (uint32_t partition = 0; partition < plan.partitions; ++partition) { + workspace.remove_consumed(stages.core(partition).path); + workspace.remove_consumed(stages.spill(partition).path); + } + partition_build_ms = elapsed_milliseconds(partition_start, std::chrono::steady_clock::now()); + + auto hierarchy_start = std::chrono::steady_clock::now(); + for (uint32_t level = 1; level <= hierarchy.max_level; ++level) { + if (global_level[level - 1]) { + build_global_level(res, + *global_vectors[level - 1], + *sidecars[level - 1], + layout, + params.metric, + plan.preferred_buffer_bytes, + io); + io.sidecar_write_bytes += checked_file_mul(global_vectors[level - 1]->count(), + global_vectors[level - 1]->header().record_size, + "upper vector bytes"); + } + sidecars[level - 1]->finalize(); + RAFT_EXPECTS(sidecars[level - 1]->count() == hierarchy.active_by_level[level - 1], + "External HNSW upper sidecar count mismatch"); + io.sidecar_write_bytes += checked_file_mul(sidecars[level - 1]->count(), + sidecars[level - 1]->header().record_size, + "upper sidecar bytes"); + io.sidecar_write_requests += sidecars[level - 1]->request_count(); + if (global_level[level - 1]) { + io.sidecar_write_requests += global_vectors[level - 1]->request_count(); + } + } + + std::vector> sidecar_readers; + sidecar_readers.reserve(hierarchy.max_level); + size_t sidecar_reader_buffer = static_cast(std::max( + uint64_t{4} << 10, + std::min(uint64_t{8} << 20, + plan.staging_buffer_bytes / std::max(1, hierarchy.max_level)))); + for (uint32_t level = 1; level <= hierarchy.max_level; ++level) { + sidecar_readers.push_back(std::make_unique( + sidecars[level - 1]->descriptor(), sidecars[level - 1]->header(), sidecar_reader_buffer)); + } + std::vector upper_record(1 + layout.M); + for (uint64_t row = 0; row < rows; ++row) { + uint32_t level = + params.hierarchy == HnswHierarchy::GPU + ? level_for_internal_id(static_cast(row), hnsw_level_seed, layout.M) + : 0; + write_hnsw_upper_node_header(output, layout, level); + for (uint32_t current = 1; current <= level; ++current) { + const std::byte* record = nullptr; + RAFT_EXPECTS(sidecar_readers[current - 1]->next(record), + "External HNSW upper sidecar ended early"); + std::memcpy(upper_record.data(), record, sidecars[current - 1]->header().record_size); + RAFT_EXPECTS(upper_record[0] == row, "External HNSW upper sidecar is not in base-ID order"); + uint32_t count = 0; + while (count < layout.M && + upper_record[count + 1] != std::numeric_limits::max()) { + uint32_t neighbor = upper_record[count + 1]; + RAFT_EXPECTS(neighbor < rows, "External HNSW upper neighbor is out of range"); + RAFT_EXPECTS(level_for_internal_id(neighbor, hnsw_level_seed, layout.M) >= current, + "External HNSW upper neighbor is not active at this level"); + ++count; + } + write_hnsw_upper_block(output, layout, upper_record.data() + 1, count); + } + } + for (auto& reader : sidecar_readers) { + io.sidecar_read_bytes += reader->bytes_read(); + io.sidecar_read_requests += reader->read_requests(); + } + + output.flush(); + RAFT_EXPECTS(output.position() == exact_output, + "External HNSW output size mismatch: expected %zu, got %zu", + static_cast(exact_output), + static_cast(output.position())); + io.output_write_bytes = output.position(); + io.output_write_requests = output.request_count(); + hierarchy_ms = elapsed_milliseconds(hierarchy_start, std::chrono::steady_clock::now()); + + auto publication_start = std::chrono::steady_clock::now(); + workspace.update_manifest("publishing", total_core, total_spill, exact_output); + workspace.publish(output_fd); + + hnsw_index->set_file_descriptor( + cuvs::util::file_descriptor(workspace.final_path().string(), O_RDONLY)); + auto build_end = std::chrono::steady_clock::now(); + publication_ms = elapsed_milliseconds(publication_start, build_end); + auto elapsed = elapsed_milliseconds(build_start, build_end); + RAFT_LOG_INFO( + "External HNSW build complete in %ld ms: preflight %ld ms, centroid training %ld ms, " + "assignment/staging %ld ms, partition processing %ld ms (loads %ld ms cumulative, CAGRA " + "builds %ld ms, base serialization %ld ms), hierarchy %ld ms, publication %ld ms. Logical " + "I/O: sample %.3f GiB, source %.3f GiB, " + "stage write/read %.3f/%.3f GiB, sidecar write/read %.3f/%.3f GiB, output %.3f GiB; " + "KvikIO requests stage write/read %zu/%zu, output %zu", + elapsed, + preflight_ms, + centroid_training_ms, + assignment_stage_ms, + partition_build_ms, + partition_load_ms, + cagra_build_ms, + base_serialization_ms, + hierarchy_ms, + publication_ms, + io.sample_read_bytes / static_cast(uint64_t{1} << 30), + io.source_read_bytes / static_cast(uint64_t{1} << 30), + io.stage_write_bytes / static_cast(uint64_t{1} << 30), + io.stage_read_bytes / static_cast(uint64_t{1} << 30), + io.sidecar_write_bytes / static_cast(uint64_t{1} << 30), + io.sidecar_read_bytes / static_cast(uint64_t{1} << 30), + io.output_write_bytes / static_cast(uint64_t{1} << 30), + static_cast(io.stage_write_requests), + static_cast(io.stage_read_requests), + static_cast(io.output_write_requests)); + return hnsw_index; + } catch (const std::exception& error) { + workspace.mark_failed("failed", error.what()); + throw; + } catch (...) { + workspace.mark_failed("failed", "unknown exception"); + throw; + } +} + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/src/neighbors/detail/hnsw/external_format.hpp b/cpp/src/neighbors/detail/hnsw/external_format.hpp new file mode 100644 index 0000000000..c16db1e501 --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/external_format.hpp @@ -0,0 +1,222 @@ +/* + * 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 + +namespace cuvs::neighbors::hnsw::detail::external { + +constexpr uint32_t stage_schema_version = 1; +constexpr uint64_t stage_data_offset = 4096; +constexpr uint64_t stage_magic = UINT64_C(0x43555653484e5357); // "CUVSHNSW" +constexpr uint32_t little_endian_marker = UINT32_C(0x01020304); + +enum class stage_kind : uint32_t { core = 1, spill = 2, upper_vectors = 3, upper_links = 4 }; +enum class element_kind : uint32_t { f32 = 1, f16 = 2, i8 = 3, u8 = 4 }; + +template +constexpr element_kind element_kind_for() +{ + if constexpr (std::is_same_v) { + return element_kind::f32; + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + return element_kind::f16; + } else if constexpr (std::is_same_v) { + return element_kind::i8; + } else if constexpr (std::is_same_v) { + return element_kind::u8; + } else { + static_assert(!sizeof(T), "unsupported external HNSW element type"); + } +} + +struct stage_header { + uint64_t magic = stage_magic; + uint32_t version = stage_schema_version; + uint32_t endian = little_endian_marker; + uint32_t kind = 0; + uint32_t element = 0; + uint32_t element_size = 0; + uint32_t dimension = 0; + uint32_t partition = 0; + uint32_t level = 0; + uint64_t record_size = 0; + uint64_t committed_records = 0; + uint64_t data_offset = stage_data_offset; + uint64_t parameter_fingerprint = 0; + std::array reserved{}; +}; +static_assert(sizeof(stage_header) == 128); +static_assert(std::is_trivially_copyable_v); + +inline uint64_t checked_file_add(uint64_t lhs, uint64_t rhs, std::string_view what) +{ + if (rhs > std::numeric_limits::max() - lhs) { + throw std::overflow_error("overflow computing " + std::string{what}); + } + return lhs + rhs; +} + +inline uint64_t checked_file_mul(uint64_t lhs, uint64_t rhs, std::string_view what) +{ + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + throw std::overflow_error("overflow computing " + std::string{what}); + } + return lhs * rhs; +} + +inline uint64_t expected_file_size(const stage_header& header) +{ + return checked_file_add( + header.data_offset, + checked_file_mul(header.committed_records, header.record_size, "stage payload size"), + "stage file size"); +} + +inline void pwrite_all(int fd, const void* source, size_t bytes, uint64_t offset) +{ + auto* ptr = static_cast(source); + size_t done = 0; + while (done < bytes) { + uint64_t current = checked_file_add(offset, done, "external HNSW write offset"); + if (current > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("external HNSW write offset exceeds off_t"); + } + ssize_t written = ::pwrite(fd, ptr + done, bytes - done, static_cast(current)); + if (written < 0) { + if (errno == EINTR) { continue; } + throw std::runtime_error("external HNSW pwrite failed: " + std::string{strerror(errno)}); + } + if (written == 0) { throw std::runtime_error("external HNSW pwrite made no progress"); } + done += static_cast(written); + } +} + +inline void pread_all(int fd, void* destination, size_t bytes, uint64_t offset) +{ + auto* ptr = static_cast(destination); + size_t done = 0; + while (done < bytes) { + uint64_t current = checked_file_add(offset, done, "external HNSW read offset"); + if (current > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("external HNSW read offset exceeds off_t"); + } + ssize_t read_bytes = ::pread(fd, ptr + done, bytes - done, static_cast(current)); + if (read_bytes < 0) { + if (errno == EINTR) { continue; } + throw std::runtime_error("external HNSW pread failed: " + std::string{strerror(errno)}); + } + if (read_bytes == 0) { throw std::runtime_error("truncated external HNSW file"); } + done += static_cast(read_bytes); + } +} + +inline void write_stage_header(int fd, const stage_header& header) +{ + static_assert(std::endian::native == std::endian::little, + "external HNSW stage format currently requires a little-endian host"); + pwrite_all(fd, &header, sizeof(header), 0); +} + +inline stage_header read_stage_header(int fd) +{ + stage_header header; + pread_all(fd, &header, sizeof(header), 0); + return header; +} + +inline void validate_stage_header(const stage_header& header, + stage_kind expected_kind, + element_kind expected_element, + uint32_t expected_element_size, + uint32_t expected_dimension, + uint32_t expected_partition, + uint32_t expected_level, + uint64_t expected_record_size, + uint64_t expected_fingerprint) +{ + RAFT_EXPECTS(header.magic == stage_magic, "Invalid external HNSW stage magic"); + RAFT_EXPECTS(header.version == stage_schema_version, + "Unsupported external HNSW stage version %u", + header.version); + RAFT_EXPECTS(header.endian == little_endian_marker, + "External HNSW stage endianness does not match this host"); + RAFT_EXPECTS(header.kind == static_cast(expected_kind), + "External HNSW stage kind mismatch"); + RAFT_EXPECTS(header.element == static_cast(expected_element) && + header.element_size == expected_element_size, + "External HNSW stage element type mismatch"); + RAFT_EXPECTS(header.dimension == expected_dimension, "External HNSW stage dimension mismatch"); + RAFT_EXPECTS(header.partition == expected_partition, "External HNSW stage partition mismatch"); + RAFT_EXPECTS(header.level == expected_level, "External HNSW stage level mismatch"); + RAFT_EXPECTS(header.record_size == expected_record_size, + "External HNSW stage record size mismatch"); + RAFT_EXPECTS(header.data_offset == stage_data_offset, "External HNSW stage data offset mismatch"); + RAFT_EXPECTS(header.parameter_fingerprint == expected_fingerprint, + "External HNSW stage parameter fingerprint mismatch"); +} + +inline void validate_stage_file_size(int fd, const stage_header& header) +{ + struct stat stat_buffer{}; + if (::fstat(fd, &stat_buffer) != 0) { + throw std::runtime_error("external HNSW fstat failed: " + std::string{strerror(errno)}); + } + RAFT_EXPECTS(stat_buffer.st_size >= 0, "Invalid external HNSW stage size"); + RAFT_EXPECTS(static_cast(stat_buffer.st_size) == expected_file_size(header), + "External HNSW stage size mismatch: expected %zu, got %zu", + static_cast(expected_file_size(header)), + static_cast(stat_buffer.st_size)); +} + +template +stage_header make_stage_header(stage_kind kind, + uint32_t dimension, + uint32_t partition, + uint32_t level, + uint64_t record_size, + uint64_t parameter_fingerprint) +{ + stage_header header; + header.kind = static_cast(kind); + header.element = static_cast(element_kind_for()); + header.element_size = sizeof(T); + header.dimension = dimension; + header.partition = partition; + header.level = level; + header.record_size = record_size; + header.parameter_fingerprint = parameter_fingerprint; + return header; +} + +inline uint64_t fnv1a64(const void* bytes, size_t size) +{ + constexpr uint64_t offset = UINT64_C(14695981039346656037); + constexpr uint64_t prime = UINT64_C(1099511628211); + uint64_t value = offset; + auto* ptr = static_cast(bytes); + for (size_t i = 0; i < size; ++i) { + value ^= ptr[i]; + value *= prime; + } + return value; +} + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/src/neighbors/detail/hnsw/external_translate.cu b/cpp/src/neighbors/detail/hnsw/external_translate.cu new file mode 100644 index 0000000000..bcc3690a0d --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/external_translate.cu @@ -0,0 +1,39 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "external_translate.hpp" + +#include +#include + +namespace cuvs::neighbors::hnsw::detail::external { +namespace { + +struct translate_graph_id { + const uint32_t* local_to_global; + size_t mapping_size; + + __device__ uint32_t operator()(uint32_t local) const + { + return local < mapping_size ? local_to_global[local] : UINT32_MAX; + } +}; + +} // namespace + +void translate_graph_ids(raft::resources const& res, + uint32_t* graph, + size_t graph_size, + const uint32_t* local_to_global, + size_t mapping_size) +{ + auto graph_view = raft::make_device_vector_view(graph, graph_size); + raft::linalg::map(res, + graph_view, + translate_graph_id{local_to_global, mapping_size}, + raft::make_const_mdspan(graph_view)); +} + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/src/neighbors/detail/hnsw/external_translate.hpp b/cpp/src/neighbors/detail/hnsw/external_translate.hpp new file mode 100644 index 0000000000..26da6f4c2e --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/external_translate.hpp @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include +#include + +namespace cuvs::neighbors::hnsw::detail::external { + +CUVS_EXPORT void translate_graph_ids(raft::resources const& res, + uint32_t* graph, + size_t graph_size, + const uint32_t* local_to_global, + size_t mapping_size); + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/src/neighbors/detail/hnsw/external_workspace.hpp b/cpp/src/neighbors/detail/hnsw/external_workspace.hpp new file mode 100644 index 0000000000..152bc78f94 --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/external_workspace.hpp @@ -0,0 +1,253 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "external_format.hpp" + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::hnsw::detail::external { + +class external_workspace { + public: + explicit external_workspace(std::filesystem::path build_dir, uint64_t parameter_fingerprint) + : build_dir_(std::move(build_dir)), fingerprint_(parameter_fingerprint) + { + RAFT_EXPECTS(!build_dir_.empty(), "ACE build_dir must not be empty"); + std::error_code error; + if (!std::filesystem::exists(build_dir_, error)) { + RAFT_EXPECTS(std::filesystem::create_directories(build_dir_, error) && !error, + "failed to create external HNSW build directory: %s", + error.message().c_str()); + } + RAFT_EXPECTS(std::filesystem::is_directory(build_dir_, error) && !error, + "external HNSW build_dir is not a directory"); + + final_path_ = build_dir_ / "hnsw_index.bin"; + RAFT_EXPECTS(!std::filesystem::exists(final_path_), + "refusing to overwrite existing HNSW index: %s", + final_path_.c_str()); + + static std::atomic sequence{0}; + for (int attempt = 0; attempt < 100; ++attempt) { + auto value = sequence.fetch_add(1, std::memory_order_relaxed); + std::ostringstream name; + name << ".hnsw-external-" << static_cast(::getpid()) << "-" << value; + invocation_dir_ = build_dir_ / name.str(); + if (::mkdir(invocation_dir_.c_str(), 0700) == 0) { break; } + if (errno != EEXIST) { + RAFT_FAIL("failed to create external HNSW workspace: %s", strerror(errno)); + } + invocation_dir_.clear(); + } + RAFT_EXPECTS(!invocation_dir_.empty(), "failed to create a unique external HNSW workspace"); + + manifest_path_ = invocation_dir_ / "manifest.json"; + manifest_fd_ = + cuvs::util::file_descriptor(manifest_path_.string(), O_CREAT | O_EXCL | O_RDWR, 0600); + partial_path_ = invocation_dir_ / "hnsw_index.bin.partial"; + update_manifest("created"); + } + + external_workspace(const external_workspace&) = delete; + external_workspace& operator=(const external_workspace&) = delete; + external_workspace(external_workspace&&) = delete; + external_workspace& operator=(external_workspace&&) = delete; + + // Always drop the private staging directory, including after a failed build. + // A successfully published hnsw_index.bin is not owned here and is left in place. + ~external_workspace() noexcept { cleanup_noexcept(); } + + [[nodiscard]] const std::filesystem::path& invocation_dir() const noexcept + { + return invocation_dir_; + } + [[nodiscard]] const std::filesystem::path& final_path() const noexcept { return final_path_; } + + [[nodiscard]] std::filesystem::path private_path(std::string_view filename) const + { + RAFT_EXPECTS( + !(filename.empty() || filename.find('/') != std::string_view::npos || + filename.find('\\') != std::string_view::npos || filename == "." || filename == ".."), + "invalid external HNSW private filename"); + return invocation_dir_ / filename; + } + + cuvs::util::file_descriptor create_private_file(std::string_view filename, mode_t mode = 0600) + { + auto path = private_path(filename); + cuvs::util::file_descriptor fd(path.string(), O_CREAT | O_EXCL | O_RDWR, mode); + owned_files_.push_back(path); + return fd; + } + + template + cuvs::util::file_descriptor create_stage_file(std::string_view filename, + stage_kind kind, + uint32_t dimension, + uint32_t partition, + uint32_t level, + uint64_t record_size) + { + auto fd = create_private_file(filename); + auto header = + make_stage_header(kind, dimension, partition, level, record_size, fingerprint_); + RAFT_EXPECTS(::ftruncate(fd.get(), static_cast(stage_data_offset)) == 0, + "failed to initialize external HNSW stage file: %s", + strerror(errno)); + write_stage_header(fd.get(), header); + return fd; + } + + cuvs::util::file_descriptor create_partial(uint64_t exact_size) + { + auto fd = create_private_file("hnsw_index.bin.partial"); + RAFT_EXPECTS(exact_size <= static_cast(std::numeric_limits::max()), + "external HNSW output exceeds off_t"); + int status = ::posix_fallocate(fd.get(), 0, static_cast(exact_size)); + RAFT_EXPECTS(status == 0, "failed to preallocate external HNSW output: %s", strerror(status)); + return fd; + } + + void update_manifest(std::string_view phase, + uint64_t core_records = 0, + uint64_t spill_records = 0, + uint64_t output_bytes = 0, + std::string_view failure_reason = {}) + { + std::ostringstream json; + json << "{\n" + << " \"schema_version\": 1,\n" + << " \"parameter_fingerprint\": " << fingerprint_ << ",\n" + << " \"phase\": \"" << json_escape(phase) << "\",\n" + << " \"core_records\": " << core_records << ",\n" + << " \"spill_records\": " << spill_records << ",\n" + << " \"output_bytes\": " << output_bytes << ",\n" + << " \"failure_reason\": \"" << json_escape(failure_reason) << "\"\n" + << "}\n"; + auto contents = json.str(); + RAFT_EXPECTS(::ftruncate(manifest_fd_.get(), 0) == 0, + "failed to truncate external HNSW manifest"); + pwrite_all(manifest_fd_.get(), contents.data(), contents.size(), 0); + RAFT_EXPECTS(::fsync(manifest_fd_.get()) == 0, "failed to sync external HNSW manifest"); + } + + // Best-effort note before the destructor removes the private workspace. + void mark_failed(std::string_view phase, std::string_view reason) noexcept + { + try { + update_manifest(phase, 0, 0, 0, reason); + } catch (...) { + } + } + + void remove_consumed(const std::filesystem::path& path) + { + auto owned = std::find(owned_files_.begin(), owned_files_.end(), path); + RAFT_EXPECTS(owned != owned_files_.end(), + "refusing to remove a file not owned by this HNSW workspace"); + std::error_code error; + RAFT_EXPECTS(std::filesystem::remove(path, error) && !error, + "failed to remove consumed external HNSW stage: %s", + error.message().c_str()); + owned_files_.erase(owned); + } + + void publish(cuvs::util::file_descriptor& partial_fd) + { + RAFT_EXPECTS( + ::fsync(partial_fd.get()) == 0, "failed to sync external HNSW output: %s", strerror(errno)); + partial_fd.close(); + + // link(2) provides no-replace atomic publication on the same filesystem. Removing the private + // name afterwards leaves the complete inode reachable only through hnsw_index.bin. + if (::link(partial_path_.c_str(), final_path_.c_str()) != 0) { + if (errno == EEXIST) { + RAFT_FAIL("refusing to overwrite an HNSW index published concurrently"); + } + RAFT_FAIL("failed to publish external HNSW index: %s", strerror(errno)); + } + if (::unlink(partial_path_.c_str()) == 0) { + auto partial = std::find(owned_files_.begin(), owned_files_.end(), partial_path_); + if (partial != owned_files_.end()) { owned_files_.erase(partial); } + } + sync_directory(build_dir_); + } + + private: + static std::string json_escape(std::string_view value) + { + std::string out; + out.reserve(value.size()); + for (char character : value) { + switch (character) { + case '\\': out += "\\\\"; break; + case '"': out += "\\\""; break; + case '\n': out += "\\n"; break; + case '\r': out += "\\r"; break; + case '\t': out += "\\t"; break; + default: out += static_cast(character) < 0x20 ? '?' : character; + } + } + return out; + } + + static void sync_directory(const std::filesystem::path& path) + { +#ifdef O_DIRECTORY + int descriptor = ::open(path.c_str(), O_RDONLY | O_DIRECTORY); +#else + int descriptor = ::open(path.c_str(), O_RDONLY); +#endif + RAFT_EXPECTS(descriptor >= 0, "failed to open directory for sync: %s", strerror(errno)); + int result = ::fsync(descriptor); + int saved = errno; + ::close(descriptor); + RAFT_EXPECTS(result == 0 || saved == EINVAL, "failed to sync directory: %s", strerror(saved)); + } + + // Drops private staging files and the invocation directory. Does not remove a published + // hnsw_index.bin or unrelated files already in build_dir_. + void cleanup_noexcept() noexcept + { + manifest_fd_.close(); + std::error_code error; + for (const auto& path : owned_files_) { + std::filesystem::remove(path, error); + error.clear(); + } + std::filesystem::remove(manifest_path_, error); + error.clear(); + std::filesystem::remove(invocation_dir_, error); + } + + std::filesystem::path build_dir_; + std::filesystem::path invocation_dir_; + std::filesystem::path manifest_path_; + std::filesystem::path partial_path_; + std::filesystem::path final_path_; + uint64_t fingerprint_; + cuvs::util::file_descriptor manifest_fd_; + std::vector owned_files_; +}; + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/src/neighbors/detail/hnsw/index_impl.hpp b/cpp/src/neighbors/detail/hnsw/index_impl.hpp new file mode 100644 index 0000000000..78b1cde52f --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/index_impl.hpp @@ -0,0 +1,174 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-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::neighbors::hnsw::detail { + +// This is needed as hnswlib hardcodes the distance type to float +// or int32_t in certain places. However, we can solve uint8 or int8 +// natively with the patch cuVS applies. We could potentially remove +// all the hardcodes and propagate templates throughout hnswlib, but +// as of now it's not needed. +template +struct hnsw_dist_t { + using type = void; +}; + +template <> +struct hnsw_dist_t { + using type = float; +}; + +template <> +struct hnsw_dist_t { + using type = float; +}; + +template <> +struct hnsw_dist_t { + using type = int; +}; + +template <> +struct hnsw_dist_t { + using type = int; +}; + +template +struct index_impl : index { + public: + /** + * @brief load a base-layer-only hnswlib index originally saved from a built CAGRA index + * + * @param[in] filepath path to the index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + * @param[in] hierarchy hierarchy used for upper HNSW layers + */ + index_impl(int dim, cuvs::distance::DistanceType metric, HnswHierarchy hierarchy) + : index{dim, metric, hierarchy} + { + if (metric == cuvs::distance::DistanceType::InnerProduct) { + space_ = std::make_unique::type>>(dim); + } else if (metric == cuvs::distance::DistanceType::L2Expanded) { + if constexpr (std::is_same_v || std::is_same_v) { + space_ = std::make_unique::type>>(dim); + } else if constexpr (std::is_same_v or std::is_same_v) { + space_ = std::make_unique>(dim); + } + } + + RAFT_EXPECTS(space_ != nullptr, "Unsupported metric type was used"); + } + + /** + @brief Get hnswlib index + */ + auto get_index() const -> void const* override { return appr_alg_.get(); } + + /** + @brief Set ef for search + */ + void set_ef(int ef) const override + { + ensure_loaded(); + appr_alg_->ef_ = ef; + } + + /** + @brief Set index + */ + void set_index(std::unique_ptr::type>>&& index) + { + appr_alg_ = std::move(index); + } + + /** + @brief Get space + */ + auto get_space() const -> hnswlib::SpaceInterface::type>* + { + return space_.get(); + } + + /** + @brief Set file descriptor for disk-backed index + */ + void set_file_descriptor(cuvs::util::file_descriptor&& fd) { hnsw_fd_.emplace(std::move(fd)); } + + /** + @brief Get file descriptor + */ + auto file_descriptor() const -> const std::optional& + { + return hnsw_fd_; + } + + /** + @brief Get file path for disk-backed index + */ + std::string file_path() const override + { + if (hnsw_fd_.has_value() && hnsw_fd_->is_valid()) { return hnsw_fd_->get_path(); } + return ""; + } + + /** + @brief Ensure the index is loaded into memory. + If the index is disk-backed and not yet loaded, this will load it from the file. + */ + void ensure_loaded() const + { + if (appr_alg_ != nullptr) { return; } // Already loaded + + // Check if we have a file descriptor to load from + if (!hnsw_fd_.has_value() || !hnsw_fd_->is_valid()) { + RAFT_FAIL("Cannot load HNSW index: no file descriptor available and index not in memory"); + } + + std::string filepath = hnsw_fd_->get_path(); + RAFT_EXPECTS(!filepath.empty(), "Cannot load HNSW index: file path is empty"); + RAFT_EXPECTS(std::filesystem::exists(filepath), + "Cannot load HNSW index: file does not exist: %s", + filepath.c_str()); + + RAFT_LOG_INFO("Loading HNSW index from disk: %s", filepath.c_str()); + + try { + appr_alg_ = std::make_unique::type>>( + space_.get(), filepath); + if (this->hierarchy() == HnswHierarchy::NONE) { appr_alg_->base_layer_only = true; } + } catch (const std::bad_alloc& e) { + RAFT_FAIL( + "Failed to load HNSW index from '%s': insufficient host memory. " + "The index is too large to fit in available RAM. " + "Consider using a machine with more memory or reducing the dataset size.", + filepath.c_str()); + } + } + + private: + mutable std::unique_ptr::type>> appr_alg_; + std::unique_ptr::type>> space_; + std::optional hnsw_fd_; +}; + +} // namespace cuvs::neighbors::hnsw::detail diff --git a/cpp/src/neighbors/detail/hnsw/serialize_layout.hpp b/cpp/src/neighbors/detail/hnsw/serialize_layout.hpp new file mode 100644 index 0000000000..7a3dfc016e --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw/serialize_layout.hpp @@ -0,0 +1,327 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "external_format.hpp" + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::hnsw::detail::external { + +constexpr size_t hnswlib_header_size = 96; +constexpr uint64_t hnsw_level_seed = 100; + +inline uint64_t splitmix64(uint64_t value) +{ + value += UINT64_C(0x9e3779b97f4a7c15); + value = (value ^ (value >> 30)) * UINT64_C(0xbf58476d1ce4e5b9); + value = (value ^ (value >> 27)) * UINT64_C(0x94d049bb133111eb); + return value ^ (value >> 31); +} + +inline uint32_t level_for_internal_id(uint32_t id, uint64_t seed, size_t M) +{ + if (M < 2) { throw std::invalid_argument("HNSW M must be at least two"); } + constexpr uint64_t mantissa_mask = (UINT64_C(1) << 53) - 1; + constexpr uint64_t denominator = (UINT64_C(1) << 53) + 1; + const uint64_t sample = + ((splitmix64(static_cast(id) ^ seed) >> 11) & mantissa_mask) + 1; + uint64_t threshold = denominator / M; + uint32_t level = 0; + while (sample <= threshold) { + ++level; + threshold /= M; + } + return level; +} + +struct hierarchy_summary { + uint32_t max_level = 0; + uint32_t entry_point = 0; + std::vector active_by_level; + uint64_t total_active_occurrences = 0; +}; + +inline hierarchy_summary summarize_hierarchy(uint64_t rows, + size_t M, + uint64_t seed = hnsw_level_seed) +{ + if (rows == 0 || rows > std::numeric_limits::max()) { + throw std::invalid_argument("invalid row count for HNSW hierarchy"); + } + hierarchy_summary summary; + for (uint64_t row = 0; row < rows; ++row) { + auto level = level_for_internal_id(static_cast(row), seed, M); + if (level > summary.max_level || + (level == summary.max_level && static_cast(row) > summary.entry_point)) { + summary.max_level = level; + summary.entry_point = static_cast(row); + } + if (summary.active_by_level.size() < level) { summary.active_by_level.resize(level, 0); } + for (uint32_t current = 0; current < level; ++current) { + ++summary.active_by_level[current]; + ++summary.total_active_occurrences; + } + } + return summary; +} + +struct hnsw_serialize_layout { + size_t offset_level0 = 0; + size_t rows = 0; + size_t size_data_per_element = 0; + size_t label_offset = 0; + size_t offset_data = 0; + int max_level = 0; + uint32_t entry_point = 0; + size_t maxM = 0; + size_t maxM0 = 0; + size_t M = 0; + double mult = 0; + size_t ef_construction = 0; + size_t dimension = 0; + size_t element_size = 0; + size_t graph_degree = 0; + size_t upper_block_size = 0; + + [[nodiscard]] uint64_t base_section_size() const + { + return checked_file_mul(rows, size_data_per_element, "HNSW base section"); + } + + [[nodiscard]] uint64_t exact_file_size(uint64_t upper_occurrences) const + { + uint64_t size = + checked_file_add(hnswlib_header_size, base_section_size(), "HNSW header and base section"); + size = checked_file_add( + size, checked_file_mul(rows, sizeof(uint32_t), "HNSW upper node headers"), "HNSW file"); + return checked_file_add( + size, + checked_file_mul(upper_occurrences, upper_block_size, "HNSW upper link blocks"), + "HNSW exact file size"); + } +}; + +template +hnsw_serialize_layout make_hnsw_serialize_layout(hnswlib::SpaceInterface* space, + uint64_t rows, + size_t dimension, + size_t graph_degree, + size_t ef_construction, + bool hierarchy, + const hierarchy_summary& summary) +{ + if (rows == 0 || rows > std::numeric_limits::max()) { + throw std::invalid_argument("invalid HNSW row count"); + } + auto algorithm = std::make_unique>( + space, 1, (graph_degree + 1) / 2, ef_construction); + hnsw_serialize_layout layout; + layout.offset_level0 = algorithm->offsetLevel0_; + layout.rows = static_cast(rows); + layout.size_data_per_element = algorithm->size_data_per_element_; + layout.label_offset = algorithm->label_offset_; + layout.offset_data = algorithm->offsetData_; + layout.max_level = hierarchy ? static_cast(summary.max_level) : 1; + layout.entry_point = hierarchy ? summary.entry_point : static_cast(rows / 2); + layout.maxM = algorithm->maxM_; + layout.maxM0 = algorithm->maxM0_; + layout.M = algorithm->M_; + layout.mult = algorithm->mult_; + layout.ef_construction = algorithm->ef_construction_; + layout.dimension = dimension; + layout.element_size = sizeof(T); + layout.graph_degree = graph_degree; + layout.upper_block_size = algorithm->size_links_per_element_; + + const size_t expected_base = + sizeof(uint32_t) + layout.maxM0 * sizeof(uint32_t) + dimension * sizeof(T) + sizeof(size_t); + RAFT_EXPECTS(layout.size_data_per_element == expected_base, "Unexpected hnswlib base row layout"); + RAFT_EXPECTS(layout.maxM0 >= graph_degree && layout.maxM0 - graph_degree <= 1, + "Unexpected hnswlib base degree"); + RAFT_EXPECTS(layout.upper_block_size == sizeof(uint32_t) + layout.M * sizeof(uint32_t), + "Unexpected hnswlib upper row layout"); + return layout; +} + +template +hnsw_serialize_layout hnsw_serialize_layout_from_algorithm( + const hnswlib::HierarchicalNSW& algorithm, + uint64_t rows, + size_t dimension, + size_t graph_degree) +{ + hnsw_serialize_layout layout; + layout.offset_level0 = algorithm.offsetLevel0_; + layout.rows = static_cast(rows); + layout.size_data_per_element = algorithm.size_data_per_element_; + layout.label_offset = algorithm.label_offset_; + layout.offset_data = algorithm.offsetData_; + layout.max_level = algorithm.maxlevel_; + layout.entry_point = algorithm.enterpoint_node_; + layout.maxM = algorithm.maxM_; + layout.maxM0 = algorithm.maxM0_; + layout.M = algorithm.M_; + layout.mult = algorithm.mult_; + layout.ef_construction = algorithm.ef_construction_; + layout.dimension = dimension; + layout.element_size = sizeof(T); + layout.graph_degree = graph_degree; + layout.upper_block_size = algorithm.size_links_per_element_; + return layout; +} + +class sequential_file_writer { + public: + sequential_file_writer(const cuvs::util::file_descriptor& fd, + uint64_t start_offset, + size_t buffer_size) + : fd_(fd), file_offset_(start_offset), buffer_(std::max(buffer_size, 1)) + { + const auto path = fd_.get_path(); + if (!path.empty()) { file_handle_ = std::make_unique(path, "r+"); } + } + + sequential_file_writer(const sequential_file_writer&) = delete; + sequential_file_writer& operator=(const sequential_file_writer&) = delete; + + ~sequential_file_writer() noexcept + { + try { + flush(); + } catch (...) { + } + } + + void write(const void* source, size_t bytes) + { + auto* source_bytes = static_cast(source); + while (bytes != 0) { + if (used_ == buffer_.size()) { flush(); } + size_t chunk = std::min(bytes, buffer_.size() - used_); + std::memcpy(buffer_.data() + used_, source_bytes, chunk); + used_ += chunk; + source_bytes += chunk; + bytes -= chunk; + } + } + + void flush() + { + if (used_ == 0) { return; } + if (file_handle_) { + const size_t written = file_handle_->pwrite(buffer_.data(), used_, file_offset_).get(); + RAFT_EXPECTS( + written == used_, "Incomplete HNSW write: expected %zu bytes, wrote %zu", used_, written); + } else { + cuvs::util::write_large_file(fd_, buffer_.data(), used_, file_offset_); + } + file_offset_ = checked_file_add(file_offset_, used_, "sequential HNSW output offset"); + used_ = 0; + ++requests_; + } + + [[nodiscard]] uint64_t position() const { return file_offset_ + used_; } + [[nodiscard]] uint64_t request_count() const noexcept { return requests_; } + + private: + const cuvs::util::file_descriptor& fd_; + uint64_t file_offset_; + std::vector buffer_; + std::unique_ptr file_handle_; + size_t used_ = 0; + uint64_t requests_ = 0; +}; + +template +void write_hnsw_header_fields(Writer& writer, const hnsw_serialize_layout& layout) +{ + writer.write(reinterpret_cast(&layout.offset_level0), sizeof(layout.offset_level0)); + writer.write(reinterpret_cast(&layout.rows), sizeof(layout.rows)); + writer.write(reinterpret_cast(&layout.rows), sizeof(layout.rows)); + writer.write(reinterpret_cast(&layout.size_data_per_element), + sizeof(layout.size_data_per_element)); + writer.write(reinterpret_cast(&layout.label_offset), sizeof(layout.label_offset)); + writer.write(reinterpret_cast(&layout.offset_data), sizeof(layout.offset_data)); + writer.write(reinterpret_cast(&layout.max_level), sizeof(layout.max_level)); + writer.write(reinterpret_cast(&layout.entry_point), sizeof(layout.entry_point)); + writer.write(reinterpret_cast(&layout.maxM), sizeof(layout.maxM)); + writer.write(reinterpret_cast(&layout.maxM0), sizeof(layout.maxM0)); + writer.write(reinterpret_cast(&layout.M), sizeof(layout.M)); + writer.write(reinterpret_cast(&layout.mult), sizeof(layout.mult)); + writer.write(reinterpret_cast(&layout.ef_construction), + sizeof(layout.ef_construction)); +} + +inline void write_hnsw_header(sequential_file_writer& writer, const hnsw_serialize_layout& layout) +{ + write_hnsw_header_fields(writer, layout); + RAFT_EXPECTS(writer.position() == hnswlib_header_size, + "HNSW header size mismatch: expected %zu, got %zu", + hnswlib_header_size, + static_cast(writer.position())); +} + +template +void write_hnsw_base_row(Writer& writer, + const hnsw_serialize_layout& layout, + const uint32_t* neighbors, + const T* vector, + uint32_t external_label) +{ + uint32_t count = static_cast(layout.graph_degree); + writer.write(reinterpret_cast(&count), sizeof(count)); + writer.write(reinterpret_cast(neighbors), layout.graph_degree * sizeof(uint32_t)); + uint32_t zero = 0; + for (size_t index = layout.graph_degree; index < layout.maxM0; ++index) { + writer.write(reinterpret_cast(&zero), sizeof(zero)); + } + writer.write(reinterpret_cast(vector), layout.dimension * sizeof(T)); + size_t label = external_label; + writer.write(reinterpret_cast(&label), sizeof(label)); +} + +template +void write_hnsw_upper_node_header(Writer& writer, + const hnsw_serialize_layout& layout, + uint32_t level) +{ + uint64_t bytes = checked_file_mul(level, layout.upper_block_size, "HNSW node upper links"); + if (bytes > std::numeric_limits::max()) { + throw std::overflow_error("HNSW node upper link list exceeds uint32_t"); + } + uint32_t link_list_size = static_cast(bytes); + writer.write(reinterpret_cast(&link_list_size), sizeof(link_list_size)); +} + +template +void write_hnsw_upper_block(Writer& writer, + const hnsw_serialize_layout& layout, + const uint32_t* neighbors, + uint32_t count) +{ + RAFT_EXPECTS(count <= layout.M, "HNSW upper neighbor count exceeds M"); + writer.write(reinterpret_cast(&count), sizeof(count)); + writer.write(reinterpret_cast(neighbors), count * sizeof(uint32_t)); + uint32_t zero = 0; + for (size_t index = count; index < layout.M; ++index) { + writer.write(reinterpret_cast(&zero), sizeof(zero)); + } +} + +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 744bc6a7a2..ef4b48c96f 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -295,6 +295,18 @@ ConfigureTest( ) if(BUILD_CAGRA_HNSWLIB) + ConfigureTest( + NAME NEIGHBORS_HNSW_EXTERNAL_IO_TEST + PATH neighbors/hnsw_external_io.cu + GPUS 1 + PERCENT 100 + ) + target_link_libraries(NEIGHBORS_HNSW_EXTERNAL_IO_TEST PRIVATE hnswlib::hnswlib kvikio::kvikio) + target_compile_definitions(NEIGHBORS_HNSW_EXTERNAL_IO_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) + target_compile_options( + NEIGHBORS_HNSW_EXTERNAL_IO_TEST PRIVATE $<$:--diag-suppress=68> + ) + ConfigureTest( NAME NEIGHBORS_HNSW_TEST PATH neighbors/hnsw.cu diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 03e9746b9b..e8f7c96391 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -218,7 +218,7 @@ void test_hnsw_ace_build_does_not_truncate_existing_index() std::string contents; index_file >> contents; EXPECT_EQ(contents, expected_contents); - EXPECT_TRUE(std::filesystem::exists(workspace.path() / "cagra_graph.npy")); + EXPECT_FALSE(std::filesystem::exists(workspace.path() / "cagra_graph.npy")); } template @@ -293,6 +293,16 @@ class AnnHnswAceTest : public ::testing::TestWithParam { hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); ASSERT_NE(hnsw_index, nullptr); + EXPECT_EQ(hnsw_index->file_path(), + (std::filesystem::path(temp_dir) / "hnsw_index.bin").string()); + EXPECT_TRUE(std::filesystem::exists(hnsw_index->file_path())); + EXPECT_FALSE(std::filesystem::exists(std::filesystem::path(temp_dir) / "cagra_graph.npy")); + EXPECT_FALSE( + std::filesystem::exists(std::filesystem::path(temp_dir) / "reordered_dataset.npy")); + EXPECT_FALSE( + std::filesystem::exists(std::filesystem::path(temp_dir) / "augmented_dataset.npy")); + EXPECT_FALSE( + std::filesystem::exists(std::filesystem::path(temp_dir) / "dataset_mapping.npy")); // Prepare queries on host auto queries_host = raft::make_host_matrix(ps.n_queries, ps.dim); @@ -307,35 +317,6 @@ class AnnHnswAceTest : public ::testing::TestWithParam { search_params.ef = std::max(ps.ef_construction, ps.k * 2); search_params.num_threads = 1; - if (!ps.use_disk) { - hnsw::search(handle_, - search_params, - *hnsw_index, - queries_host.view(), - indexes_hnsw_host.view(), - distances_hnsw_host.view()); - for (size_t i = 0; i < queries_size; i++) { - indexes_hnsw[i] = indexes_hnsw_host.data_handle()[i]; - distances_hnsw[i] = distances_hnsw_host.data_handle()[i]; - } - - // Convert indexes for comparison - std::vector indexes_hnsw_converted(queries_size); - for (size_t i = 0; i < queries_size; i++) { - indexes_hnsw_converted[i] = static_cast(indexes_hnsw[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)) - << "HNSW ACE build and search failed recall check"; - } - tmp_index_file index_file; hnsw::serialize(handle_, index_file.filename, *hnsw_index); @@ -385,7 +366,8 @@ class AnnHnswAceTest : public ::testing::TestWithParam { void testHnswAceMemoryLimitFallback() { - // This test verifies that setting tiny memory limits forces disk mode automatically + // A positive memory cap is a hard planner limit. A cap below the minimum partition peak must + // fail before any staging artifact is created. // Create temporary directory for ACE build std::string temp_dir = std::string("/tmp/cuvs_hnsw_ace_memlimit_test_") + std::to_string(std::time(nullptr)) + "_" + @@ -404,32 +386,21 @@ class AnnHnswAceTest : public ::testing::TestWithParam { hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; hnsw_params.M = 32; - // Configure ACE parameters with tiny memory limits to force disk mode auto ace_params = graph_build_params::ace_params(); ace_params.npartitions = ps.npartitions; ace_params.ef_construction = ps.ef_construction; ace_params.build_dir = temp_dir; - ace_params.use_disk = false; // Not explicitly requesting disk mode + ace_params.use_disk = false; - // Selected host memory limit should enforce disk mode ace_params.max_host_memory_gb = 0.001; ace_params.max_gpu_memory_gb = 3.0; hnsw_params.graph_build_params = ace_params; - // Build HNSW index using ACE - should automatically fall back to disk mode - auto hnsw_index = - hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); - - ASSERT_NE(hnsw_index, nullptr); - - // Verify that disk mode was triggered by checking for the expected files - std::string graph_file = temp_dir + "/cagra_graph.npy"; - std::string reordered_file = temp_dir + "/reordered_dataset.npy"; - - EXPECT_TRUE(std::filesystem::exists(graph_file)) - << "Graph file should exist when memory limit triggers disk mode fallback"; - EXPECT_TRUE(std::filesystem::exists(reordered_file)) - << "Reordered dataset file should exist when memory limit triggers disk mode fallback"; + EXPECT_THROW(hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())), + raft::logic_error); + EXPECT_FALSE(std::filesystem::exists(temp_dir + "/hnsw_index.bin")); + EXPECT_FALSE(std::filesystem::exists(temp_dir + "/cagra_graph.npy")); + EXPECT_FALSE(std::filesystem::exists(temp_dir + "/reordered_dataset.npy")); } // Clean up temporary directory @@ -684,22 +655,21 @@ inline std::vector generate_hnsw_ace_inputs() ); } -// Inputs specifically for testing memory limit fallback to disk mode +// Inputs for a host/GPU cap below the minimum planned partition peak. inline std::vector generate_hnsw_ace_memory_fallback_inputs() { return { - // Test with L2 metric {10, // n_queries 5000, // n_rows 64, // dim 10, // k 2, // npartitions 100, // ef_construction - false, // use_disk (not explicitly set, should be triggered by memory limit) + false, // use_disk (unused: the cap fails the planner before staging) cuvs::distance::DistanceType::L2Expanded, - 0.0, // min_recall (not checked in fallback test) - 0.001, // max_host_memory_gb (tiny limit to force disk mode) - 0.001} // max_gpu_memory_gb (tiny limit to force disk mode) + 0.0, // min_recall (not checked) + 0.001, // max_host_memory_gb (below the minimum planned partition peak) + 0.001} // max_gpu_memory_gb }; } diff --git a/cpp/tests/neighbors/hnsw_external_io.cu b/cpp/tests/neighbors/hnsw_external_io.cu new file mode 100644 index 0000000000..a622133e38 --- /dev/null +++ b/cpp/tests/neighbors/hnsw_external_io.cu @@ -0,0 +1,458 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../src/neighbors/detail/cagra/ace_external_plan.hpp" +#include "../../src/neighbors/detail/hnsw.hpp" +#include "../../src/neighbors/detail/hnsw/external_format.hpp" +#include "../../src/neighbors/detail/hnsw/external_translate.hpp" +#include "../../src/neighbors/detail/hnsw/external_workspace.hpp" +#include "../../src/neighbors/detail/hnsw/serialize_layout.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::hnsw::detail::external { +namespace { + +class temporary_directory { + public: + temporary_directory() + { + static uint64_t sequence = 0; + do { + path_ = std::filesystem::temp_directory_path() / + ("cuvs-hnsw-external-test-" + std::to_string(::getpid()) + "-" + + std::to_string(sequence++)); + } while (!std::filesystem::create_directory(path_)); + } + ~temporary_directory() + { + std::error_code error; + std::filesystem::remove_all(path_, error); + } + const std::filesystem::path& path() const noexcept { return path_; } + + private: + std::filesystem::path path_; +}; + +TEST(HnswExternalPlan, BigannByteLedger) +{ + cagra::detail::ace_external_plan_input input; + input.rows = 1'000'000'000; + input.dim = 128; + input.element_size = sizeof(uint8_t); + input.M = 24; + input.graph_degree = 48; + input.intermediate_degree = 72; + input.requested_partitions = 64; + input.available_host_bytes = uint64_t{2} << 40; + input.available_device_bytes = uint64_t{2} << 40; + input.force_disk = true; + input.hierarchy = true; + + auto plan = cagra::detail::make_ace_external_plan(input); + auto materialized_host = + cagra::detail::estimate_materialized_cagra_ace_host_bytes(input, plan.partitions); + constexpr double gib = static_cast(uint64_t{1} << 30); + EXPECT_TRUE(plan.use_disk); + EXPECT_GT(materialized_host, input.rows * input.graph_degree * sizeof(uint32_t)); + EXPECT_NEAR((plan.bytes.source_scan + plan.bytes.centroid_sample) / gib, 120.4, 0.2); + EXPECT_NEAR((plan.bytes.stage_write + plan.bytes.stage_read) / gib, 506.6, 0.2); + EXPECT_NEAR(plan.bytes.base_output / gib, 309.2, 0.2); + EXPECT_NEAR((plan.bytes.upper_sidecar_write + plan.bytes.upper_sidecar_read) / gib, 8.1, 0.2); + EXPECT_NEAR(plan.bytes.final_upper_output / gib, 7.8, 0.2); + EXPECT_NEAR(plan.bytes.logical_total() / gib, 952.1, 0.5); +} + +TEST(HnswHostMemory, EstimateIncludesRuntimeAndHierarchyStorage) +{ + constexpr int64_t rows = 1'000; + constexpr int64_t dim = 128; + constexpr int degree = 48; + const auto base_file_bytes = + rows * (sizeof(uint32_t) + degree * sizeof(uint32_t) + dim * sizeof(float) + sizeof(size_t)); + auto base_only = cuvs::neighbors::hnsw::detail::estimate_hnsw_host_memory( + rows, dim, degree, cuvs::distance::DistanceType::L2Expanded, HnswHierarchy::NONE, 200); + auto with_hierarchy = cuvs::neighbors::hnsw::detail::estimate_hnsw_host_memory( + rows, dim, degree, cuvs::distance::DistanceType::L2Expanded, HnswHierarchy::GPU, 200); + EXPECT_GT(base_only, base_file_bytes); + EXPECT_GT(with_hierarchy, base_only); +} + +TEST(HnswExternalPlan, OverflowAndHardCap) +{ + cagra::detail::ace_external_plan_input input; + input.rows = std::numeric_limits::max(); + input.dim = std::numeric_limits::max(); + input.element_size = sizeof(float); + input.M = 32; + input.graph_degree = 64; + input.intermediate_degree = 96; + input.available_host_bytes = uint64_t{1} << 30; + input.available_device_bytes = uint64_t{1} << 30; + input.force_disk = true; + EXPECT_THROW(cagra::detail::make_ace_external_plan(input), raft::logic_error); + + EXPECT_THROW(cagra::detail::external_checked_mul( + std::numeric_limits::max(), uint64_t{2}, "test overflow"), + raft::logic_error); + + input.rows = 10'000; + input.dim = 128; + input.available_host_bytes = 1; + input.available_device_bytes = 1; + EXPECT_THROW(cagra::detail::make_ace_external_plan(input), raft::logic_error); +} + +TEST(HnswExternalPlan, PartitionsIncreaseMonotonically) +{ + cagra::detail::ace_external_plan_input input; + input.rows = 1'000'000; + input.dim = 128; + input.element_size = sizeof(float); + input.M = 24; + input.graph_degree = 48; + input.intermediate_degree = 72; + input.available_host_bytes = uint64_t{8} << 30; + input.available_device_bytes = uint64_t{8} << 30; + input.force_disk = true; + auto roomy = cagra::detail::make_ace_external_plan(input); + input.available_host_bytes = uint64_t{2} << 30; + input.available_device_bytes = uint64_t{2} << 30; + auto tight = cagra::detail::make_ace_external_plan(input); + EXPECT_GE(tight.partitions, roomy.partitions); + EXPECT_EQ(roomy.queue_depth, 2); + EXPECT_LE(tight.queue_depth, roomy.queue_depth); + EXPECT_LE(tight.host_peak_bytes, input.available_host_bytes * 4 / 5); + EXPECT_LE(tight.device_peak_bytes, input.available_device_bytes * 4 / 5); + + input.requested_partitions = input.rows; + input.available_host_bytes = uint64_t{64} << 30; + input.available_device_bytes = uint64_t{64} << 30; + auto clamped = cagra::detail::make_ace_external_plan(input); + EXPECT_LE(clamped.partitions, + cagra::detail::external_maximum_partitions(input.rows, input.intermediate_degree)); + EXPECT_GT(clamped.target_occurrences, input.intermediate_degree); + + input.requested_partitions = input.rows + 1; + EXPECT_THROW(cagra::detail::make_ace_external_plan(input), raft::logic_error); +} + +TEST(HnswExternalPlan, SampleAndAssignmentRangesAreBoundedAndMonotonic) +{ + auto sample = cagra::detail::make_external_sample_ranges(10'003, 1'001); + ASSERT_LE(sample.size(), 16); + uint64_t sampled = 0; + uint64_t previous_end = 0; + for (const auto& range : sample) { + EXPECT_GT(range.count, 0); + EXPECT_GE(range.start, previous_end); + EXPECT_LE(range.start + range.count, 10'003); + sampled += range.count; + previous_end = range.start + range.count; + } + EXPECT_EQ(sampled, 1'001); + + auto scan = cagra::detail::make_external_monotonic_ranges(10'003, 257); + uint64_t scanned = 0; + for (const auto& range : scan) { + EXPECT_EQ(range.start, scanned); + EXPECT_GT(range.count, 0); + scanned += range.count; + } + EXPECT_EQ(scanned, 10'003); + + auto single = cagra::detail::make_external_sample_ranges(100, 1); + ASSERT_EQ(single.size(), 1); + EXPECT_EQ(single[0].start, 49); + EXPECT_EQ(single[0].count, 1); +} + +TEST(HnswExternalConversion, CopiesAndConvertsRows) +{ + std::array float_source{1.0f, -2.0f, 3.5f, 4.0f}; + std::array float_destination{}; + copy_rows_as_float(float_source.data(), float_destination.data(), 2, 2); + EXPECT_EQ(float_destination, float_source); + + std::array int_source{1, -2, 3, 4}; + std::array converted{}; + copy_rows_as_float(int_source.data(), converted.data(), 1, int_source.size()); + EXPECT_EQ(converted, (std::array{1.0f, -2.0f, 3.0f, 4.0f})); + + EXPECT_THROW( + copy_rows_as_float(float_source.data(), float_destination.data(), uint64_t{1} << 63, 2), + std::overflow_error); +} + +TEST(HnswExternalFormat, RoundTripAndTruncation) +{ + temporary_directory directory; + auto path = directory.path() / "core.stage"; + cuvs::util::file_descriptor fd(path.string(), O_CREAT | O_EXCL | O_RDWR, 0600); + auto header = make_stage_header( + stage_kind::core, 3, 7, 0, 2 * sizeof(uint32_t) + 3 * sizeof(float), 123); + header.committed_records = 1; + ASSERT_EQ(::ftruncate(fd.get(), static_cast(expected_file_size(header))), 0); + write_stage_header(fd.get(), header); + std::array record{}; + uint32_t label = 9; + std::memcpy(record.data(), &label, sizeof(label)); + pwrite_all(fd.get(), record.data(), record.size(), stage_data_offset); + + auto decoded = read_stage_header(fd.get()); + EXPECT_NO_THROW(validate_stage_header( + decoded, stage_kind::core, element_kind::f32, sizeof(float), 3, 7, 0, record.size(), 123)); + EXPECT_NO_THROW(validate_stage_file_size(fd.get(), decoded)); + ASSERT_EQ(::ftruncate(fd.get(), static_cast(expected_file_size(header) - 1)), 0); + EXPECT_THROW(validate_stage_file_size(fd.get(), decoded), raft::logic_error); +} + +TEST(HnswExternalFormat, ShortReadIsCatchable) +{ + temporary_directory directory; + auto path = directory.path() / "short.bin"; + cuvs::util::file_descriptor fd(path.string(), O_CREAT | O_EXCL | O_RDWR, 0600); + uint8_t byte = 1; + pwrite_all(fd.get(), &byte, sizeof(byte), 0); + std::array output{}; + EXPECT_THROW(pread_all(fd.get(), output.data(), output.size(), 0), std::runtime_error); +} + +TEST(HnswExternalFormat, WriteFailureIsCatchable) +{ + if (!std::filesystem::exists("/dev/full")) { GTEST_SKIP() << "/dev/full is unavailable"; } + cuvs::util::file_descriptor fd("/dev/full", O_WRONLY); + uint8_t byte = 1; + EXPECT_THROW(pwrite_all(fd.get(), &byte, sizeof(byte), 0), std::runtime_error); +} + +TEST(HnswExternalStage, RoundRobinAssignmentsRemainBuffered) +{ + temporary_directory directory; + external_workspace workspace(directory.path(), 23); + constexpr uint32_t partitions = 16; + constexpr uint32_t dimension = 4; + stage_buffer_pool stages( + workspace, dimension, partitions, uint64_t{64} << 10, uint64_t{64} << 10, 23); + std::array vector{}; + for (uint32_t row = 0; row < 100; ++row) { + uint32_t partition = row % partitions; + stages.append_core(partition, row, vector.data()); + stages.append_spill((partition + 1) % partitions, partition, row / partitions, vector.data()); + } + stages.finalize(); + EXPECT_LE(stages.write_requests(), 2 * partitions) + << "round-robin assignment should retain one buffer per stage file"; +} + +TEST(HnswExternalStage, BufferedReaderReusesFileAcrossRefills) +{ + temporary_directory directory; + external_workspace workspace(directory.path(), 29); + constexpr uint32_t dimension = 4; + stage_buffer_pool stages( + workspace, dimension, 1, uint64_t{64} << 10, uint64_t{64} << 10, 29); + std::array vector{}; + for (uint32_t row = 0; row < 3; ++row) { + stages.append_core(0, row, vector.data()); + } + stages.finalize(); + + cuvs::util::file_descriptor fd(stages.core(0).path.string(), O_RDONLY); + auto header = read_stage_header(fd.get()); + buffered_stage_reader reader(fd, header, static_cast(header.record_size)); + const std::byte* record = nullptr; + for (uint32_t row = 0; row < 3; ++row) { + ASSERT_TRUE(reader.next(record)); + uint32_t label = 0; + std::memcpy(&label, record, sizeof(label)); + EXPECT_EQ(label, row); + } + EXPECT_FALSE(reader.next(record)); + EXPECT_EQ(reader.read_requests(), 3); +} + +TEST(HnswExternalGraph, OwnerTranslationIsPartitionLocal) +{ + std::vector prefixes{0, 3, 7, 10}; + EXPECT_EQ(external_owner_id(prefixes, 0, 2), 2); + EXPECT_EQ(external_owner_id(prefixes, 1, 3), 6); + EXPECT_EQ(external_owner_id(prefixes, 2, 2), 9); + EXPECT_THROW(external_owner_id(prefixes, 1, 4), raft::logic_error); + EXPECT_THROW(external_owner_id(prefixes, 3, 0), raft::logic_error); +} + +TEST(HnswExternalGraph, TranslatesIdsOnDevice) +{ + raft::resources res; + std::array graph{0, 2, 1, 3, 7}; + std::array mapping{10, 20, 30, 40}; + auto device_graph = raft::make_device_vector(res, graph.size()); + auto device_mapping = raft::make_device_vector(res, mapping.size()); + raft::copy(res, + device_graph.view(), + raft::make_host_vector_view(graph.data(), graph.size())); + raft::copy(res, + device_mapping.view(), + raft::make_host_vector_view(mapping.data(), mapping.size())); + translate_graph_ids(res, + device_graph.data_handle(), + device_graph.size(), + device_mapping.data_handle(), + device_mapping.size()); + raft::copy(res, + raft::make_host_vector_view(graph.data(), graph.size()), + raft::make_const_mdspan(device_graph.view())); + raft::resource::sync_stream(res); + EXPECT_EQ(graph, (std::array{10, 30, 20, 40, UINT32_MAX})); +} + +TEST(HnswExternalWorkspace, PreservesCallerFilesAndPublishesAtomically) +{ + temporary_directory directory; + auto sentinel = directory.path() / "sentinel.txt"; + { + std::ofstream stream(sentinel); + stream << "caller-owned"; + } + + std::filesystem::path failed_workspace; + { + external_workspace workspace(directory.path(), 17); + failed_workspace = workspace.invocation_dir(); + auto stage = workspace.create_private_file("owned.stage"); + uint8_t value = 7; + pwrite_all(stage.get(), &value, 1, 0); + workspace.mark_failed("test", "injected"); + } + EXPECT_TRUE(std::filesystem::exists(sentinel)); + EXPECT_FALSE(std::filesystem::exists(failed_workspace)); + + { + external_workspace workspace(directory.path(), 18); + auto partial = workspace.create_partial(4); + uint32_t value = 42; + pwrite_all(partial.get(), &value, sizeof(value), 0); + workspace.publish(partial); + } + EXPECT_TRUE(std::filesystem::exists(directory.path() / "hnsw_index.bin")); + EXPECT_TRUE(std::filesystem::exists(sentinel)); + EXPECT_THROW(external_workspace(directory.path(), 19), raft::logic_error); +} + +TEST(HnswExternalHierarchy, DeterministicNestedDistribution) +{ + constexpr uint32_t rows = 1'000'000; + uint64_t level_one = 0; + uint64_t level_two = 0; + for (uint32_t id = 0; id < rows; ++id) { + auto first = level_for_internal_id(id, hnsw_level_seed, 24); + auto second = level_for_internal_id(id, hnsw_level_seed, 24); + EXPECT_EQ(first, second); + level_one += first >= 1; + level_two += first >= 2; + } + EXPECT_NEAR(static_cast(level_one) / rows, 1.0 / 24.0, 0.001); + EXPECT_NEAR(static_cast(level_two) / rows, 1.0 / (24.0 * 24.0), 0.0002); + EXPECT_LE(level_two, level_one); + auto summary = summarize_hierarchy(rows, 24); + ASSERT_FALSE(summary.active_by_level.empty()); + EXPECT_EQ(summary.active_by_level[0], level_one); +} + +TEST(HnswExternalLayout, WritesLoadableNoneHierarchy) +{ + temporary_directory directory; + auto path = directory.path() / "tiny.bin"; + auto batched_path = directory.path() / "tiny-batched.bin"; + cuvs::util::file_descriptor fd(path.string(), O_CREAT | O_EXCL | O_RDWR, 0600); + hnswlib::L2Space space(2); + hierarchy_summary summary; + auto layout = make_hnsw_serialize_layout(&space, 4, 2, 2, 100, false, summary); + auto exact = layout.exact_file_size(0); + ASSERT_EQ(::posix_fallocate(fd.get(), 0, static_cast(exact)), 0); + sequential_file_writer writer(fd, 0, 8); + write_hnsw_header(writer, layout); + std::array, 4> vectors{{{0, 0}, {1, 0}, {0, 1}, {1, 1}}}; + std::array, 4> neighbors{{{1, 2}, {0, 3}, {0, 3}, {1, 2}}}; + for (uint32_t row = 0; row < 4; ++row) { + write_hnsw_base_row(writer, layout, neighbors[row].data(), vectors[row].data(), row); + } + for (uint32_t row = 0; row < 4; ++row) { + write_hnsw_upper_node_header(writer, layout, 0); + } + writer.flush(); + EXPECT_EQ(writer.position(), exact); + EXPECT_GT(writer.request_count(), 1); + fd.close(); + + raft::resources res; + index_params params; + params.hierarchy = HnswHierarchy::NONE; + params.ef_construction = 100; + { + std::ofstream batched(batched_path, std::ios::binary | std::ios::trunc); + ASSERT_TRUE(batched); + serialize_to_hnswlib_batched( + res, + batched, + params, + vectors.size(), + vectors[0].size(), + neighbors[0].size(), + cuvs::distance::DistanceType::L2Expanded, + [&](int64_t start, int64_t count, auto graph_batch, auto dataset_batch, auto label_batch) { + for (int64_t row = 0; row < count; ++row) { + for (int64_t edge = 0; edge < static_cast(neighbors[0].size()); ++edge) { + graph_batch(row, edge) = neighbors[start + row][edge]; + } + for (int64_t column = 0; column < static_cast(vectors[0].size()); ++column) { + dataset_batch(row, column) = vectors[start + row][column]; + } + label_batch(row) = static_cast(start + row); + } + }); + } + std::ifstream shared_bytes(path, std::ios::binary); + std::ifstream batched_bytes(batched_path, std::ios::binary); + ASSERT_TRUE(shared_bytes && batched_bytes); + std::vector shared_contents((std::istreambuf_iterator(shared_bytes)), + std::istreambuf_iterator()); + std::vector batched_contents((std::istreambuf_iterator(batched_bytes)), + std::istreambuf_iterator()); + EXPECT_EQ(shared_contents, batched_contents); + + EXPECT_NO_THROW({ + hnswlib::HierarchicalNSW loaded(&space, path.string()); + EXPECT_EQ(loaded.getCurrentElementCount(), 4); + }); + EXPECT_NO_THROW({ + hnswlib::HierarchicalNSW loaded(&space, batched_path.string()); + EXPECT_EQ(loaded.getCurrentElementCount(), 4); + }); +} + +} // namespace +} // namespace cuvs::neighbors::hnsw::detail::external diff --git a/examples/cpp/src/hnsw_ace_example.cu b/examples/cpp/src/hnsw_ace_example.cu index 22354193e5..2107df399a 100644 --- a/examples/cpp/src/hnsw_ace_example.cu +++ b/examples/cpp/src/hnsw_ace_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -47,12 +47,11 @@ void hnsw_build_search_ace(raft::device_resources const& dev_resources, // sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the // partition sizes (up to 3x in our tests). ace_params.npartitions = 4; - // Set the directory to store the ACE build artifacts. This should be the fastest disk in the - // system and hold enough space for twice the dataset, final graph, and label mapping. - ace_params.build_dir = "/tmp/hnsw_ace_build"; - // Set whether to use disk-based storage for ACE build. When true, enables disk-based operations - // for memory-efficient graph construction. If not set, the index will be built in memory if the - // graph fits in host and GPU memory, and on disk otherwise. + // Directory for the final HNSW index and private partition staging. Use fast disk with room + // for the index plus the staging high-water mark. ACE writes hnsw_index.bin here and does not + // write a full CAGRA graph or dataset mapping. Failed builds remove the private staging + // directory and do not leave hnsw_index.bin. + ace_params.build_dir = "/tmp/hnsw_ace_build"; ace_params.use_disk = true; hnsw_params.graph_build_params = ace_params; // Set M parameter to control the graph degree (graph_degree = m * 2, intermediate_graph_degree = diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py index 183d530e7c..5d6ac601d2 100644 --- a/python/cuvs/cuvs/tests/test_hnsw_ace.py +++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -28,11 +28,8 @@ def run_hnsw_ace_build_search_test( expected_recall=0.9, ): """ - Test HNSW index build using ACE algorithm. - - - Build HNSW index using ACE via hnsw.build() - - For disk mode: serialize -> deserialize -> search - - For in-memory mode: search directly + Test HNSW index build using ACE via hnsw.build(). + ACE writes hnsw_index.bin; search uses the deserialized index. """ dataset = generate_data((n_rows, n_cols), dtype) queries = generate_data((n_queries, n_cols), dtype) @@ -67,37 +64,30 @@ def run_hnsw_ace_build_search_test( hnsw_index = hnsw.build(index_params, dataset) assert hnsw_index.trained + hnsw_file = os.path.join(temp_dir, "hnsw_index.bin") + assert os.path.exists(hnsw_file) + for cagra_artifact in ( + "cagra_graph.npy", + "reordered_dataset.npy", + "augmented_dataset.npy", + "dataset_mapping.npy", + ): + assert not os.path.exists(os.path.join(temp_dir, cagra_artifact)) + + deserialized_index = hnsw.load( + index_params, + hnsw_file, + n_cols, + dtype, + metric=metric, + ) - if use_disk: - # For disk mode, the index is serialized to disk by the build function - # We need to deserialize it before searching - hnsw_file = os.path.join(temp_dir, "hnsw_index.bin") - assert os.path.exists(hnsw_file) - - # Deserialize from disk for searching - deserialized_index = hnsw.load( - index_params, - hnsw_file, - n_cols, - dtype, - metric=metric, - ) - - # Search the deserialized index - search_params = hnsw.SearchParams( - ef=max(ef_construction, k * 2), num_threads=1 - ) - out_dist, out_idx = hnsw.search( - search_params, deserialized_index, queries, k - ) - else: - # For in-memory mode, search directly - search_params = hnsw.SearchParams( - ef=max(ef_construction, k * 2), num_threads=1 - ) - out_dist, out_idx = hnsw.search( - search_params, hnsw_index, queries, k - ) + search_params = hnsw.SearchParams( + ef=max(ef_construction, k * 2), num_threads=1 + ) + out_dist, out_idx = hnsw.search( + search_params, deserialized_index, queries, k + ) # Calculate reference values with sklearn skl_metric = { @@ -223,8 +213,8 @@ def test_hnsw_ace_disk_serialize_deserialize(): assert recall >= 0.7, f"Recall {recall:.3f} is below expected 0.7" -def test_hnsw_ace_tiny_memory_limit_triggers_disk_mode(): - """Test that setting tiny memory limits triggers disk mode automatically.""" +def test_hnsw_ace_tiny_memory_limit_fails_before_staging(): + """A host/GPU cap below the minimum planned partition peak fails before staging.""" n_rows = 5000 n_cols = 64 dtype = np.float32 @@ -233,14 +223,13 @@ def test_hnsw_ace_tiny_memory_limit_triggers_disk_mode(): dataset = generate_data((n_rows, n_cols), dtype) with tempfile.TemporaryDirectory() as temp_dir: - # Set ACE parameters with memory limits slightly above the minimum required - # This should force disk mode even though we didn't explicitly set use_disk=True + # Cap below the minimum planned partition peak. ace_params = hnsw.AceParams( npartitions=2, build_dir=temp_dir, - use_disk=False, # Not explicitly requesting disk mode - max_host_memory_gb=0.001, # Tiny limit to force disk mode - max_gpu_memory_gb=0.0, # No GPU memory limit + use_disk=False, + max_host_memory_gb=0.001, # Below the minimum planned partition peak + max_gpu_memory_gb=0.0, ) # Create HNSW index params with ACE @@ -252,18 +241,85 @@ def test_hnsw_ace_tiny_memory_limit_triggers_disk_mode(): ace_params=ace_params, ) - # Build the index using ACE - should automatically use disk mode - hnsw_index = hnsw.build(index_params, dataset) - assert hnsw_index.trained + with pytest.raises(RuntimeError, match="cap|memory|peak"): + hnsw.build(index_params, dataset) - # In disk mode, the graph should be stored in the build directory - # Check that the graph file was created - graph_file = os.path.join(temp_dir, "cagra_graph.npy") - reordered_file = os.path.join(temp_dir, "reordered_dataset.npy") + assert not os.path.exists(os.path.join(temp_dir, "hnsw_index.bin")) + assert not os.path.exists(os.path.join(temp_dir, "cagra_graph.npy")) + assert not os.path.exists( + os.path.join(temp_dir, "reordered_dataset.npy") + ) + assert not os.path.exists( + os.path.join(temp_dir, "augmented_dataset.npy") + ) - assert os.path.exists(graph_file), ( - "Graph file should exist when disk mode is triggered" + +def test_hnsw_ace_memmap_below_dataset_host_cap(): + """Build from a read-only memmap larger than the configured host cap.""" + n_rows = 20_000 + n_cols = 256 + n_queries = 32 + k = 10 + + with tempfile.TemporaryDirectory() as temp_dir: + dataset_path = os.path.join(temp_dir, "dataset.dat") + writable = np.memmap( + dataset_path, + dtype=np.float32, + mode="w+", + shape=(n_rows, n_cols), + ) + writable[:] = generate_data((n_rows, n_cols), np.float32) + writable.flush() + del writable + dataset = np.memmap( + dataset_path, + dtype=np.float32, + mode="r", + shape=(n_rows, n_cols), + ) + assert dataset.nbytes > 18 * 1024**2 + queries = np.asarray(dataset[:n_queries]).copy() + + ace_params = hnsw.AceParams( + npartitions=32, + build_dir=temp_dir, + use_disk=True, + max_host_memory_gb=18 / 1024, ) - assert os.path.exists(reordered_file), ( - "Reordered dataset file should exist when disk mode is triggered" + index_params = hnsw.IndexParams( + hierarchy="none", + M=16, + ef_construction=100, + metric="sqeuclidean", + ace_params=ace_params, + ) + built = hnsw.build(index_params, dataset) + index_path = os.path.join(temp_dir, "hnsw_index.bin") + assert built.trained + assert os.path.exists(index_path) + + loaded = hnsw.load( + index_params, + index_path, + n_cols, + np.float32, + metric="sqeuclidean", + ) + distances, neighbors = hnsw.search( + hnsw.SearchParams(ef=200, num_threads=1), + loaded, + queries, + k, + ) + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + assert ( + np.mean( + np.any( + neighbors == np.arange(n_queries, dtype=np.int64)[:, None], + axis=1, + ) + ) + >= 0.9 )