From 354c19e42c8a63b5796f7c72bd83dcdcb9049e7d Mon Sep 17 00:00:00 2001
From: cqy123456 <39671710+cqy123456@users.noreply.github.com>
Date: Wed, 8 Feb 2023 19:56:27 +0800
Subject: [PATCH] use threaddata scratch as aio and vector buffer (#667)

Signed-off-by: cqy123456 <qianya.cheng@zilliz.com>
---
 thirdparty/DiskANN/include/utils.h        |  7 ++++---
 thirdparty/DiskANN/src/pq_flash_index.cpp | 21 ++++++++++-----------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/thirdparty/DiskANN/include/utils.h b/thirdparty/DiskANN/include/utils.h
index ea7b4251a..a6fd6d002 100644
--- a/thirdparty/DiskANN/include/utils.h
+++ b/thirdparty/DiskANN/include/utils.h
@@ -203,9 +203,10 @@ namespace diskann {
 #else
     ::_aligned_free(ptr);
 #endif
-    if (malloc_trim(0) == 0) {
-      LOG(DEBUG) << "Failed to release free memory from the heap.";
-    }
+    // Note: malloc_trim() combines with jemalloc/tcmalloc may cause crash
+    // if (malloc_trim(0) == 0) {
+    //   LOG(DEBUG) << "Failed to release free memory from the heap.";
+    // }
   }
 
   inline void GenRandom(std::mt19937& rng, unsigned* addr, unsigned size,
diff --git a/thirdparty/DiskANN/src/pq_flash_index.cpp b/thirdparty/DiskANN/src/pq_flash_index.cpp
index 477a8b91f..fd75e3575 100644
--- a/thirdparty/DiskANN/src/pq_flash_index.cpp
+++ b/thirdparty/DiskANN/src/pq_flash_index.cpp
@@ -470,34 +470,36 @@ namespace diskann {
                   num_medoids * aligned_dim * sizeof(float), 32);
     std::memset(centroid_data, 0, num_medoids * aligned_dim * sizeof(float));
 
-    // borrow ctx
     ThreadData<T> data = this->thread_data.pop();
+
     while (data.scratch.sector_scratch == nullptr) {
       this->thread_data.wait_for_push_notify();
       data = this->thread_data.pop();
     }
+    // borrow ctx
     IOContext &ctx = data.ctx;
+    // borrow buf
+    auto scratch = &(data.scratch);
+    scratch->reset();
+    char *sector_scratch = scratch->sector_scratch;
+    T* medoid_coords = scratch->coord_scratch;
+
     LOG(INFO) << "Loading centroid data from medoids vector data of "
               << num_medoids << " medoid(s)";
     for (uint64_t cur_m = 0; cur_m < num_medoids; cur_m++) {
       auto medoid = medoids[cur_m];
       // read medoid nhood
-      char *medoid_buf = nullptr;
-      alloc_aligned((void **) &medoid_buf, read_len_for_node, SECTOR_LEN);
       std::vector<AlignedRead> medoid_read(1);
       medoid_read[0].len = read_len_for_node;
-      medoid_read[0].buf = medoid_buf;
+      medoid_read[0].buf = sector_scratch;
       medoid_read[0].offset = get_node_sector_offset(medoid);
       reader->read(medoid_read, ctx);
-
       // all data about medoid
-      char *medoid_node_buf = get_offset_to_node(medoid_buf, medoid);
+      char *medoid_node_buf = get_offset_to_node(sector_scratch, medoid);
 
       // add medoid coords to `coord_cache`
-      T *medoid_coords = new T[data_dim];
       T *medoid_disk_coords = OFFSET_TO_NODE_COORDS(medoid_node_buf);
       memcpy(medoid_coords, medoid_disk_coords, disk_bytes_per_point);
-
       if (!use_disk_index_pq) {
         for (uint32_t i = 0; i < data_dim; i++)
           centroid_data[cur_m * aligned_dim + i] = medoid_coords[i];
@@ -505,9 +507,6 @@ namespace diskann {
         disk_pq_table.inflate_vector((_u8 *) medoid_coords,
                                      (centroid_data + cur_m * aligned_dim));
       }
-
-      aligned_free(medoid_buf);
-      delete[] medoid_coords;
     }
 
     // return ctx