Merge branch 'track-memory' into 'mqdb-dev'

Track memory in MyScale See merge request mqdb/faiss!30
myscale · Oct 12, 2023 · c51d22d · c51d22d
2 parents ad6a732 + 869b1f0
commit c51d22d
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 1 deletion.
diff --git a/faiss/IndexFlatCodes.h b/faiss/IndexFlatCodes.h
@@ -12,6 +12,9 @@
 #include <faiss/Index.h>
 #include <faiss/impl/DistanceComputer.h>
 #include <vector>
+#ifdef MYSCALE_MODE
+#include <Common/AllocatorWithMemoryTracking.h>
+#endif
 
 namespace faiss {
 
@@ -21,7 +24,11 @@ struct IndexFlatCodes : Index {
     size_t code_size;
 
     /// encoded dataset, size ntotal * code_size
+#ifdef MYSCALE_MODE
+    std::vector<uint8_t, AllocatorWithMemoryTracking<uint8_t>> codes;
+#else
     std::vector<uint8_t> codes;
+#endif
 
     IndexFlatCodes();
 

diff --git a/faiss/impl/HNSW.h b/faiss/impl/HNSW.h
@@ -21,6 +21,10 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
 
+#ifdef MYSCALE_MODE
+#include <Common/AllocatorWithMemoryTracking.h>
+#endif
+
 #define ENABLE_STAT
 #ifdef ENABLE_STAT
 #define IF_STATISTIC(stat, inc_stat) \
@@ -113,22 +117,42 @@ struct HNSW {
     };
 
     /// assignment probability to each layer (sum=1)
+#ifdef MYSCALE_MODE
+    std::vector<double, AllocatorWithMemoryTracking<double>> assign_probas;
+#else
     std::vector<double> assign_probas;
+#endif
 
     /// number of neighbors stored per layer (cumulative), should not
     /// be changed after first add
+#ifdef MYSCALE_MODE
+    std::vector<int, AllocatorWithMemoryTracking<int>> cum_nneighbor_per_level;
+#else
     std::vector<int> cum_nneighbor_per_level;
+#endif
 
     /// level of each vector (base level = 1), size = ntotal
+#ifdef MYSCALE_MODE
+    std::vector<int, AllocatorWithMemoryTracking<int>> levels;
+#else
     std::vector<int> levels;
+#endif
 
     /// offsets[i] is the offset in the neighbors array where vector i is stored
     /// size ntotal + 1
+#ifdef MYSCALE_MODE
+    std::vector<size_t, AllocatorWithMemoryTracking<size_t>> offsets;
+#else
     std::vector<size_t> offsets;
+#endif
 
     /// neighbors[offsets[i]:offsets[i+1]] is the list of neighbors of vector i
     /// for all levels. this is where all storage goes.
+#ifdef MYSCALE_MODE
+    std::vector<storage_idx_t, AllocatorWithMemoryTracking<storage_idx_t>> neighbors;
+#else
     std::vector<storage_idx_t> neighbors;
+#endif
 
     /// entry point in the search structure (one of the points with maximum
     /// level

diff --git a/faiss/impl/HNSWfast.h b/faiss/impl/HNSWfast.h
@@ -17,6 +17,9 @@
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/random.h>
 #include <SearchIndex/VectorIndex.h>
+#ifdef MYSCALE_MODE
+#include <Common/AllocatorWithMemoryTracking.h>
+#endif
 
 namespace faiss {
 
@@ -160,7 +163,11 @@ class HNSWfast {
     };
 
     /// level of each vector (base level = 1), size = ntotal
+#ifdef MYSCALE_MODE
+    std::vector<int, AllocatorWithMemoryTracking<int>> levels;
+#else
     std::vector<int> levels;
+#endif
 
     /// entry point in the search structure (one of the points with maximum
     /// level
@@ -172,10 +179,16 @@ class HNSWfast {
     /// maximum level
     int max_level;
     int M;
+#ifdef MYSCALE_MODE
+    std::vector<char, AllocatorWithMemoryTracking<char>> level0_storage;
+    std::vector<std::vector<char, AllocatorWithMemoryTracking<char>>> link_storage;
+    std::vector<char*, AllocatorWithMemoryTracking<char*>> linkLists;
+#else
     std::vector<char> level0_storage;
     std::vector<std::vector<char>> link_storage;
-    char* level0_links;
     std::vector<char*> linkLists;
+#endif
+    char* level0_links;
     size_t level0_link_size;
     size_t link_size;
     double level_constant;