memory mapping the index more responsibly

c_api/index_io_c.h

@@ -34,6 +34,8 @@ int faiss_write_index_fname(const FaissIndex* idx, const char* fname);
 
 #define FAISS_IO_FLAG_MMAP 1
 #define FAISS_IO_FLAG_READ_ONLY 2
+#define FAISS_IO_FLAG_READ_MMAP 32
+#define FAISS_IO_FLAG_ONDISK_IVF 0x646f0000
 
 /** Read index from a file.
  * This is equivalent to `faiss:read_index` when a file descriptor is given.

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,39 @@ 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.
+
+TODO: need more asserts
+*/
+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)
+    // the size value can change for indexes with encodings like
+    // SQ, PQ etc. TODO: need to explore those.
+    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");
+
+    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 +571,17 @@ 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);
+            // TODO: have an assert around codes_ptr
+        } 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,8 @@ 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;
+// read the index from an already mmap'd data buffer
+const int IO_FLAG_READ_MMAP = 32 | 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

@@ -353,6 +353,7 @@ OnDiskInvertedLists::OnDiskInvertedLists(
           filename(filename),
           totsize(0),
           ptr(nullptr),
+          pre_mapped(false),
           read_only(false),
           locks(new LockLevels()),
           pf(new OngoingPrefetch(this)),
@@ -369,11 +370,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 +747,38 @@ InvertedLists* OnDiskInvertedListsIOHook::read(IOReader* f, int io_flags)
     return od;
 }
 
+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 unmapping and mapping responsiblity is that of the parent
+    // layer.
+    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);
+    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);
+    }
+
+    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 {
@@ -754,6 +788,10 @@ InvertedLists* OnDiskInvertedListsIOHook::read_ArrayInvertedLists(
     ails->read_only = true;
     ails->lists.resize(nlist);
 
+    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");
     FILE* fdesc = reader->f;

faiss/invlists/OnDiskInvertedLists.h

@@ -77,6 +77,7 @@ 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
 
     OnDiskInvertedLists(size_t nlist, size_t code_size, const char* filename);