MB-59575: Reducing copy overhead of already memory mapped content (#17)

- Currently, from the parent layers (zapx) we are passing a pointer to the serialized byte array in the index file, which is mmap'd on golang side.
- However, while reading the index from buffer we'd be reading (and performing memcpy's) over the entire set of vector codes in the index which can hamper performance and rss usage.
- The PR intends to avoid these memcpy's and just have a pointer to the mmap'd byte array when feasible and retrieve the specific vector codes in the read-only paths in a lazy fashion

Author: Thejas-bhat

c_api/index_io_c_ex.h

@@ -21,6 +21,14 @@
 extern "C" {
 #endif
 
+// skip prefetch phase while searching over the inverted lists
+#define FAISS_IO_FLAG_SKIP_PREFETCH 32
+// the following two macros together decide whether to read the index from an
+// already mmap'd data buffer. it's C equivalent of IO_FLAG_READ_MMAP from index_io.h
+// usage is - FAISS_IO_FLAG_READ_MMAP | FAISS_IO_FLAG_ONDISK_IVF
+#define FAISS_IO_FLAG_READ_MMAP 64
+#define FAISS_IO_FLAG_ONDISK_IVF 0x646f0000
+
 /** Write index to buffer
  */
 int faiss_write_index_buf(const FaissIndex* idx, size_t* buf_size, unsigned char** buf);

faiss/IndexFlat.cpp

@@ -169,6 +169,10 @@ FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
 }
 
 void IndexFlat::reconstruct(idx_t key, float* recons) const {
+    if (mmaped) {
+        memcpy(recons, &(codes_ptr[key * code_size]), code_size);
+        return;
+    }
     memcpy(recons, &(codes[key * code_size]), code_size);
 }
 

faiss/IndexFlat.h

@@ -54,9 +54,15 @@ struct IndexFlat : IndexFlatCodes {
 
     // get pointer to the floating point data
     float* get_xb() {
+        if (mmaped) {
+            return (float*)(codes_ptr);
+        }
         return (float*)codes.data();
     }
     const float* get_xb() const {
+        if (mmaped) {
+            return (const float*)(codes_ptr);
+        }
         return (const float*)codes.data();
     }
 

faiss/IndexFlatCodes.cpp

@@ -18,7 +18,20 @@ namespace faiss {
 IndexFlatCodes::IndexFlatCodes(size_t code_size, idx_t d, MetricType metric)
         : Index(d, metric), code_size(code_size) {}
 
-IndexFlatCodes::IndexFlatCodes() : code_size(0) {}
+IndexFlatCodes::IndexFlatCodes() :
+            code_size(0),
+            mmaped_size(0),
+            mmaped(false),
+            codes_ptr(nullptr) {}
+
+IndexFlatCodes::~IndexFlatCodes() {
+    // setting the pointer to nullptr so that the mmap'd region is zero counted
+    // from faiss side and safe to be free'd/GC'd etc. on calling application layer
+    // of faiss.
+    if (mmaped) {
+        codes_ptr = nullptr;
+    }
+}
 
 void IndexFlatCodes::add(idx_t n, const float* x) {
     FAISS_THROW_IF_NOT(is_trained);

faiss/IndexFlatCodes.h

@@ -24,9 +24,14 @@ struct IndexFlatCodes : Index {
 
     /// encoded dataset, size ntotal * code_size
     std::vector<uint8_t> codes;
+    uint8_t* codes_ptr;
+    bool mmaped; // true if codes_ptr is pointing to a mmaped region
+    size_t mmaped_size;
 
     IndexFlatCodes();
 
+    ~IndexFlatCodes() override;
+
     IndexFlatCodes(size_t code_size, idx_t d, MetricType metric = METRIC_L2);
 
     /// default add uses sa_encode

faiss/impl/index_read.cpp

@@ -188,12 +188,15 @@ static void read_ArrayInvertedLists_sizes(
 InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
     uint32_t h;
     READ1(h);
+    bool load_mem = !((io_flags & IO_FLAG_READ_MMAP) ||
+                      (io_flags & IO_FLAG_SKIP_IVF_DATA));
+
     if (h == fourcc("il00")) {
         fprintf(stderr,
                 "read_InvertedLists:"
                 " WARN! inverted lists not stored with IVF object\n");
         return nullptr;
-    } else if (h == fourcc("ilar") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
+    } else if (h == fourcc("ilar") && load_mem) {
         size_t nlist, code_size;
         READ1(nlist);
         READ1(code_size);
@@ -212,7 +215,7 @@ InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
         }
         return ails;
 
-    } else if (h == fourcc("ilar") && (io_flags & IO_FLAG_SKIP_IVF_DATA)) {
+    } else if (h == fourcc("ilar") && !load_mem) {
         // code is always ilxx where xx is specific to the type of invlists we
         // want so we get the 16 high bits from the io_flag and the 16 low bits
         // as "il"
@@ -520,6 +523,40 @@ static IndexIVFPQ* read_ivfpq(IOReader* f, uint32_t h, int io_flags) {
 
 int read_old_fmt_hack = 0;
 
+/**
+ * flat indexes which store the codes directly, can use this API instead to have a
+ * pointer to the mmaped region to avoid allocation costs. works specifically with
+ * BufIOReader as of now.
+**/
+void read_codes_mmaped(IOReader* f, IndexFlat* idxf) {
+    idxf->mmaped = true;
+
+    // read the size of codes data
+    size_t size;
+    READANDCHECK(&size, 1);
+    FAISS_THROW_IF_NOT(size >= 0 && size < (uint64_t{1} << 40));
+    size *= 4;
+
+    // size == ntotal * code_size == ntotal * d * sizeof(float) for IndexFlat
+    // NOTE: the code_size value can change for indexes with encodings like
+    // SQ, PQ although the size value will still be equal to ntotal * code_size
+    // bytes which is accessible via codes_ptr.
+    FAISS_THROW_IF_NOT(size == idxf->ntotal * idxf->code_size);
+    idxf->mmaped_size = size;
+
+    // BufIOReader is the reader which has a direct pointer to the mmaped
+    // byte array, so we can directly set the codes_ptr to the mmaped region
+    BufIOReader* reader = dynamic_cast<BufIOReader*>(f);
+    FAISS_THROW_IF_NOT_MSG(reader, "reading over mmap'd region is supported only with BufIOReader");
+    FAISS_THROW_IF_NOT_MSG(reader->buf, "reader buffer is null");
+
+    idxf->codes_ptr = const_cast<uint8_t*>(reader->buf);
+    // seek to the point where the codes section begins
+    idxf->codes_ptr += reader->rp;
+    // update read pointer appropriately
+    reader->rp += size;
+}
+
 Index* read_index(IOReader* f, int io_flags) {
     Index* idx = nullptr;
     uint32_t h;
@@ -535,9 +572,14 @@ Index* read_index(IOReader* f, int io_flags) {
         }
         read_index_header(idxf, f);
         idxf->code_size = idxf->d * sizeof(float);
-        READXBVECTOR(idxf->codes);
-        FAISS_THROW_IF_NOT(
+
+        if (io_flags & IO_FLAG_READ_MMAP) {
+            read_codes_mmaped(f, idxf);
+        } else {
+            READXBVECTOR(idxf->codes);
+            FAISS_THROW_IF_NOT(
                 idxf->codes.size() == idxf->ntotal * idxf->code_size);
+        }
         // leak!
         idx = idxf;
     } else if (h == fourcc("IxHE") || h == fourcc("IxHe")) {

faiss/index_io.h

@@ -52,6 +52,10 @@ const int IO_FLAG_ONDISK_SAME_DIR = 4;
 const int IO_FLAG_SKIP_IVF_DATA = 8;
 // don't initialize precomputed table after loading
 const int IO_FLAG_SKIP_PRECOMPUTE_TABLE = 16;
+// skip prefetch phase while searching over the inverted lists
+const int IO_FLAG_SKIP_PREFETCH = 32;
+// read the index from an already mmap'd data buffer
+const int IO_FLAG_READ_MMAP = 64 | 0x646f0000;
 // try to memmap data (useful to load an ArrayInvertedLists as an
 // OnDiskInvertedLists)
 const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000;

faiss/invlists/OnDiskInvertedLists.cpp

@@ -257,6 +257,12 @@ struct OnDiskInvertedLists::OngoingPrefetch {
 int OnDiskInvertedLists::OngoingPrefetch::global_cs = 0;
 
 void OnDiskInvertedLists::prefetch_lists(const idx_t* list_nos, int n) const {
+
+    // avoid prefetch when the ondisk-ivf is already prepared for read-only paths
+    // helpful when the queries are not batched
+    if (skip_prefetch) {
+        return;
+    }
     pf->prefetch_lists(list_nos, n);
 }
 
@@ -353,6 +359,8 @@ OnDiskInvertedLists::OnDiskInvertedLists(
           filename(filename),
           totsize(0),
           ptr(nullptr),
+          pre_mapped(false),
+          skip_prefetch(false),
           read_only(false),
           locks(new LockLevels()),
           pf(new OngoingPrefetch(this)),
@@ -369,11 +377,16 @@ OnDiskInvertedLists::~OnDiskInvertedLists() {
 
     // unmap all lists
     if (ptr != nullptr) {
-        int err = munmap(ptr, totsize);
-        if (err != 0) {
-            fprintf(stderr, "mumap error: %s", strerror(errno));
+        if (!pre_mapped) {
+            int err = munmap(ptr, totsize);
+            if (err != 0) {
+                fprintf(stderr, "mumap error: %s", strerror(errno));
+            }
+        } else {
+            ptr = nullptr;
         }
     }
+
     delete locks;
 }
 
@@ -741,10 +754,48 @@ InvertedLists* OnDiskInvertedListsIOHook::read(IOReader* f, int io_flags)
     return od;
 }
 
+/**
+ * This function is just an alternate way to use the OnDiskInvertedLists.
+ * It's useful when the index is read using BufIOReader from a uint8_t* buffer
+ * which is already mmap'd by the application layer.
+ * All the responbility of handling this mmap pointer now falls on the app layer
+**/
+InvertedLists* read_ArrayInvertedLists_MMAP(
+        IOReader* f,
+        OnDiskInvertedLists* ails,
+        const std::vector<size_t>& sizes) {
+
+    // setting this true is to ensure that the destructor does not unmap
+    // since the mmap control is on the parent layer of faiss.
+    ails->pre_mapped = true;
+
+    BufIOReader* reader = dynamic_cast<BufIOReader*>(f);
+    FAISS_THROW_IF_NOT_MSG(reader, "reading over mmap'd region is supported only with BufIOReader");
+
+    size_t o = reader->rp;
+    ails->totsize = reader->buf_size;
+    FAISS_THROW_IF_NOT(o <= ails->totsize);
+    FAISS_THROW_IF_NOT_MSG(reader->buf, "reader buffer is null");
+    // using the base pointer to the mmap'd region
+    ails->ptr = const_cast<uint8_t*>(reader->buf);
+
+    for (size_t i = 0; i < ails->nlist; i++) {
+        OnDiskInvertedLists::List& l = ails->lists[i];
+        l.size = l.capacity = sizes[i];
+        l.offset = o;
+        o += l.size * (sizeof(idx_t) + ails->code_size);
+    }
+
+    // updating the read pointer appropriately, this is needed when the IVF
+    // wrapped with another index class.
+    reader->rp = o;
+    return ails;
+}
+
 /** read from a ArrayInvertedLists into this invertedlist type */
 InvertedLists* OnDiskInvertedListsIOHook::read_ArrayInvertedLists(
         IOReader* f,
-        int /* io_flags */,
+        int io_flags,
         size_t nlist,
         size_t code_size,
         const std::vector<size_t>& sizes) const {
@@ -753,6 +804,13 @@ InvertedLists* OnDiskInvertedListsIOHook::read_ArrayInvertedLists(
     ails->code_size = code_size;
     ails->read_only = true;
     ails->lists.resize(nlist);
+    if (io_flags & IO_FLAG_SKIP_PREFETCH) {
+        ails->skip_prefetch = true;
+    }
+
+    if (io_flags & IO_FLAG_READ_MMAP) {
+        return read_ArrayInvertedLists_MMAP(f, ails, sizes);
+    }
 
     FileIOReader* reader = dynamic_cast<FileIOReader*>(f);
     FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects");

faiss/invlists/OnDiskInvertedLists.h

@@ -77,6 +77,8 @@ struct OnDiskInvertedLists : InvertedLists {
     size_t totsize;
     uint8_t* ptr;   // mmap base pointer
     bool read_only; /// are inverted lists mapped read-only
+    bool pre_mapped;// whether the content is already mmap'd before class creation
+    bool skip_prefetch; // whether to skip prefetching the lists while performing search
 
     OnDiskInvertedLists(size_t nlist, size_t code_size, const char* filename);