perf(diskann): async I/O overlap and dynamic beam

src/core/algorithm/diskann/diskann_file_reader.cc

@@ -338,6 +338,108 @@ int LinuxAlignedFileReader::read(std::vector<AlignedRead> &read_reqs,
   return ret;
 }
 
+#if (defined(__linux) || defined(__linux__))
+int LinuxAlignedFileReader::submit(PendingBatch &batch,
+                                   std::vector<AlignedRead> &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<uint32_t> &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<io_event_t> 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

src/core/algorithm/diskann/diskann_file_reader.h

@@ -58,6 +58,16 @@ struct AlignedRead {
   }
 };
 
+#if (defined(__linux) || defined(__linux__))
+struct PendingBatch {
+  std::vector<struct iocb> cbs;
+  std::vector<struct iocb *> cb_ptrs;
+  uint32_t n_submitted{0};
+  uint32_t n_reaped{0};
+  bool used_pread{false};
+};
+#endif
+
 class AlignedFileReader {
  protected:
   std::map<std::thread::id, IOContext> ctx_map;
@@ -77,6 +87,15 @@ class AlignedFileReader {
 
   virtual int read(std::vector<AlignedRead> &read_reqs, IOContext &ctx,
                    bool async = false) = 0;
+
+#if (defined(__linux) || defined(__linux__))
+  virtual int submit(PendingBatch &batch, std::vector<AlignedRead> &read_reqs,
+                     IOContext &ctx) = 0;
+
+  virtual int get_completed(PendingBatch &batch, IOContext &ctx,
+                            int min_completed,
+                            std::vector<uint32_t> &completed_indices) = 0;
+#endif
 };
 
 class LinuxAlignedFileReader : public AlignedFileReader {
@@ -101,6 +120,14 @@ class LinuxAlignedFileReader : public AlignedFileReader {
 
   int read(std::vector<AlignedRead> &read_reqs, IOContext &ctx,
            bool async = false);
+
+#if (defined(__linux) || defined(__linux__))
+  int submit(PendingBatch &batch, std::vector<AlignedRead> &read_reqs,
+             IOContext &ctx);
+
+  int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed,
+                    std::vector<uint32_t> &completed_indices);
+#endif
 };
 
 }  // namespace core

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<diskann_id_t> frontier;
-  frontier.reserve(2 * beam_width_);
+  frontier.reserve(2 * effective_beam_width);
 
   std::vector<std::pair<diskann_id_t, uint8_t *>> frontier_neighbors;
-  frontier_neighbors.reserve(2 * beam_width_);
+  frontier_neighbors.reserve(2 * effective_beam_width);
 
   std::vector<AlignedRead> frontier_read_reqs;
-  frontier_read_reqs.reserve(2 * beam_width_);
+  frontier_read_reqs.reserve(2 * effective_beam_width);
 
   std::vector<std::tuple<diskann_id_t, uint32_t, diskann_id_t *>>
       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<uint32_t> 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<diskann_id_t *>(node_buf + 1);
+          void *node_fp_coords = node_disk_buf;
 
-      cpu_timer.reset();
-      std::vector<float> 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<diskann_id_t *>(node_buf + 1);
+
+          cpu_timer.reset();
+          std::vector<float> 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();
     }
   }