From 746d7ad02a38cafb64e87c23993734431e4591b2 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 23 Jul 2026 13:53:31 +0800 Subject: [PATCH] perf: io --- .../algorithm/diskann/diskann_file_reader.cc | 102 +++++++++++++++ .../algorithm/diskann/diskann_file_reader.h | 27 ++++ src/core/algorithm/diskann/diskann_indexer.cc | 120 +++++++++++------- 3 files changed, 200 insertions(+), 49 deletions(-) diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 85be8c334..1ff8e8018 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -338,6 +338,108 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } +#if (defined(__linux) || defined(__linux__)) +int LinuxAlignedFileReader::submit(PendingBatch &batch, + std::vector &read_reqs, + IOContext &ctx) { + batch.n_submitted = 0; + batch.n_reaped = 0; + batch.used_pread = false; + batch.cbs.clear(); + batch.cb_ptrs.clear(); + + if (this->file_desc == -1) { + LOG_ERROR("submit: invalid file descriptor"); + return IndexError_Runtime; + } + + if (read_reqs.empty()) { + return 0; + } + + // If no async I/O backend is available, use synchronous pread. + if (ailego::IOBackend::Instance().is_pread()) { + int pread_ret = execute_io_pread(this->file_desc, read_reqs); + if (pread_ret != 0) { + return pread_ret; + } + batch.used_pread = true; + batch.n_submitted = (uint32_t)read_reqs.size(); + return 0; + } + + uint32_t n_ops = (uint32_t)read_reqs.size(); + batch.cbs.resize(n_ops); + batch.cb_ptrs.resize(n_ops); + + for (uint32_t j = 0; j < n_ops; j++) { + io_prep_pread(&batch.cbs[j], this->file_desc, read_reqs[j].buf, + read_reqs[j].len, read_reqs[j].offset); + batch.cbs[j].data = (void *)(uintptr_t)j; + batch.cb_ptrs[j] = &batch.cbs[j]; + } + + int ret = LibAioLoader::Instance().io_submit(ctx, (int64_t)n_ops, + batch.cb_ptrs.data()); + if (ret == (int)n_ops) { + batch.n_submitted = n_ops; + return 0; + } + + LOG_WARN("submit: io_submit returned %d (expected %u), falling back to pread", + ret, n_ops); + int pread_ret = execute_io_pread(this->file_desc, read_reqs); + if (pread_ret != 0) { + return pread_ret; + } + batch.used_pread = true; + batch.n_submitted = n_ops; + return 0; +} + +int LinuxAlignedFileReader::get_completed( + PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) { + completed_indices.clear(); + + if (batch.n_reaped >= batch.n_submitted) { + return 0; + } + + if (batch.used_pread) { + for (uint32_t i = batch.n_reaped; i < batch.n_submitted; i++) { + completed_indices.push_back(i); + } + batch.n_reaped = batch.n_submitted; + return (int)completed_indices.size(); + } + + uint32_t n_remaining = batch.n_submitted - batch.n_reaped; + int min_req = std::min((int)n_remaining, min_completed); + if (min_req < 1) min_req = 1; + + std::vector evts(n_remaining); + int ret = LibAioLoader::Instance().io_getevents( + ctx, (int64_t)min_req, (int64_t)n_remaining, evts.data(), nullptr); + if (ret < 0) { + LOG_ERROR("get_completed: io_getevents failed, ret=%d", ret); + return IndexError_Runtime; + } + + for (int i = 0; i < ret; i++) { + uint32_t idx = (uint32_t)(uintptr_t)evts[i].data; + if ((int64_t)evts[i].res != (int64_t)batch.cbs[idx].u.c.nbytes) { + LOG_WARN("get_completed: read %u failed: res=%ld, expected=%ld", idx, + (long)evts[i].res, (long)batch.cbs[idx].u.c.nbytes); + } + completed_indices.push_back(idx); + } + + batch.n_reaped += (uint32_t)ret; + return ret; +} +#endif + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..a8e1354fd 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -58,6 +58,16 @@ struct AlignedRead { } }; +#if (defined(__linux) || defined(__linux__)) +struct PendingBatch { + std::vector cbs; + std::vector cb_ptrs; + uint32_t n_submitted{0}; + uint32_t n_reaped{0}; + bool used_pread{false}; +}; +#endif + class AlignedFileReader { protected: std::map ctx_map; @@ -77,6 +87,15 @@ class AlignedFileReader { virtual int read(std::vector &read_reqs, IOContext &ctx, bool async = false) = 0; + +#if (defined(__linux) || defined(__linux__)) + virtual int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) = 0; + + virtual int get_completed(PendingBatch &batch, IOContext &ctx, + int min_completed, + std::vector &completed_indices) = 0; +#endif }; class LinuxAlignedFileReader : public AlignedFileReader { @@ -101,6 +120,14 @@ class LinuxAlignedFileReader : public AlignedFileReader { int read(std::vector &read_reqs, IOContext &ctx, bool async = false); + +#if (defined(__linux) || defined(__linux__)) + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx); + + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices); +#endif }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..62c92d589 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -791,18 +791,23 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { uint32_t num_ios = 0; + uint32_t effective_beam_width = + std::max(8u, std::min(ctx->list_size() / 5, 32u)); + std::vector frontier; - frontier.reserve(2 * beam_width_); + frontier.reserve(2 * effective_beam_width); std::vector> frontier_neighbors; - frontier_neighbors.reserve(2 * beam_width_); + frontier_neighbors.reserve(2 * effective_beam_width); std::vector frontier_read_reqs; - frontier_read_reqs.reserve(2 * beam_width_); + frontier_read_reqs.reserve(2 * effective_beam_width); std::vector> cached_neighbors; - cached_neighbors.reserve(2 * beam_width_); + cached_neighbors.reserve(2 * effective_beam_width); + + PendingBatch pending; while (candidates.has_unexpanded_node() && num_ios < io_limit_) { frontier.clear(); @@ -813,8 +818,9 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { uint64_t sector_buffer_idx = 0; uint32_t num_seen = 0; - while (candidates.has_unexpanded_node() && frontier.size() < beam_width_ && - num_seen < beam_width_) { + while (candidates.has_unexpanded_node() && + frontier.size() < effective_beam_width && + num_seen < effective_beam_width) { auto neighbor = candidates.closest_unexpanded(); num_seen++; @@ -857,11 +863,10 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { } io_timer.reset(); - - int read_ret = reader_->read(frontier_read_reqs, io_ctx); + int submit_ret = reader_->submit(pending, frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); - if (read_ret != 0) { - LOG_ERROR("cached_beam_search: reader_->read failed, ret=%d", read_ret); + if (submit_ret != 0) { + LOG_ERROR("cached_beam_search: submit failed, ret=%d", submit_ret); ctx->set_error(true); return IndexError_Runtime; } @@ -895,55 +900,72 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { for (uint64_t m = 0; m < neighbor_num; ++m) { diskann_id_t id = node_neighbors[m]; - visit_filter.set_visited(id); - - Neighbor nn(id, distances[m]); - candidates.insert(nn); + if (!visit_filter.visited(id)) { + visit_filter.set_visited(id); + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } } } - for (auto &frontier_neighbor : frontier_neighbors) { - uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( - node_per_sector_, max_node_size_, frontier_neighbor.second, - frontier_neighbor.first); - uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( - node_disk_buf, meta_.element_size()); - uint32_t neighbor_num = *node_buf; - - void *node_fp_coords = node_disk_buf; - - float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); + if (!frontier.empty()) { + std::vector completed; + while (pending.n_reaped < pending.n_submitted) { + completed.clear(); + io_timer.reset(); + int n = reader_->get_completed(pending, io_ctx, 1, completed); + stats.io_us += io_timer.micro_seconds(); + if (n < 0) { + LOG_ERROR("cached_beam_search: get_completed failed, ret=%d", n); + ctx->set_error(true); + return IndexError_Runtime; + } - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(frontier_neighbor.first))) { - topk_heap.emplace( - frontier_neighbor.first, - VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords))); - } + for (uint32_t idx : completed) { + auto &frontier_neighbor = frontier_neighbors[idx]; + uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( + node_per_sector_, max_node_size_, frontier_neighbor.second, + frontier_neighbor.first); + uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( + node_disk_buf, meta_.element_size()); + uint32_t neighbor_num = *node_buf; - diskann_id_t *node_neighbors = - reinterpret_cast(node_buf + 1); + void *node_fp_coords = node_disk_buf; - cpu_timer.reset(); - std::vector distances(neighbor_num); - pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, - ctx->pq_table_dist_buffer(), - ctx->pq_coord_buffer(), distances.data()); + float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); - stats.dist_num += neighbor_num; - stats.cpu_us += cpu_timer.micro_seconds(); + if (!ctx->filter().is_valid() || + !ctx->filter()(get_key(frontier_neighbor.first))) { + topk_heap.emplace(frontier_neighbor.first, + VectorInfo(cur_expanded_dist, + make_vector_copy(node_fp_coords))); + } - cpu_timer.reset(); - for (uint64_t m = 0; m < neighbor_num; ++m) { - diskann_id_t id = node_neighbors[m]; - visit_filter.set_visited(id); - stats.dist_num++; + diskann_id_t *node_neighbors = + reinterpret_cast(node_buf + 1); + + cpu_timer.reset(); + std::vector distances(neighbor_num); + pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, + ctx->pq_table_dist_buffer(), + ctx->pq_coord_buffer(), distances.data()); + + stats.dist_num += neighbor_num; + stats.cpu_us += cpu_timer.micro_seconds(); + + for (uint64_t m = 0; m < neighbor_num; ++m) { + diskann_id_t id = node_neighbors[m]; + if (!visit_filter.visited(id)) { + visit_filter.set_visited(id); + stats.dist_num++; + Neighbor nn(id, distances[m]); + candidates.insert(nn); + } + } - Neighbor nn(id, distances[m]); - candidates.insert(nn); + stats.cpu_us += cpu_timer.micro_seconds(); + } } - - stats.cpu_us += cpu_timer.micro_seconds(); } }