Skip to content

Commit

Permalink
Track memory in MyScale
Browse files Browse the repository at this point in the history
lqhl committed Sep 27, 2023
1 parent ad6a732 commit 869b1f0
Showing 3 changed files with 45 additions and 1 deletion.
7 changes: 7 additions & 0 deletions faiss/IndexFlatCodes.h
Original file line number Diff line number Diff line change
@@ -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();

24 changes: 24 additions & 0 deletions faiss/impl/HNSW.h
Original file line number Diff line number Diff line change
@@ -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
15 changes: 14 additions & 1 deletion faiss/impl/HNSWfast.h
Original file line number Diff line number Diff line change
@@ -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;

0 comments on commit 869b1f0

Please sign in to comment.