diff --git a/c_api/index_io_c.h b/c_api/index_io_c.h index 8935de0219..00e49b6476 100644 --- a/c_api/index_io_c.h +++ b/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. diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index 0fa3b82062..5e795facee 100644 --- a/faiss/IndexFlat.cpp +++ b/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); } diff --git a/faiss/IndexFlat.h b/faiss/IndexFlat.h index ef9910edb6..6782d1ed29 100644 --- a/faiss/IndexFlat.h +++ b/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(); } diff --git a/faiss/IndexFlatCodes.cpp b/faiss/IndexFlatCodes.cpp index eb52c76922..ab4760f02a 100644 --- a/faiss/IndexFlatCodes.cpp +++ b/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); diff --git a/faiss/IndexFlatCodes.h b/faiss/IndexFlatCodes.h index 677f3eead4..6e23f88d30 100644 --- a/faiss/IndexFlatCodes.h +++ b/faiss/IndexFlatCodes.h @@ -24,9 +24,14 @@ struct IndexFlatCodes : Index { /// encoded dataset, size ntotal * code_size std::vector 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 diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index cb668d602b..58e13f908b 100644 --- a/faiss/impl/index_read.cpp +++ b/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(f); + FAISS_THROW_IF_NOT_MSG(reader, "reading over mmap'd region is supported only with BufIOReader"); + + idxf->codes_ptr = const_cast(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")) { diff --git a/faiss/index_io.h b/faiss/index_io.h index 8d52ee1afd..29843d986c 100644 --- a/faiss/index_io.h +++ b/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; diff --git a/faiss/invlists/OnDiskInvertedLists.cpp b/faiss/invlists/OnDiskInvertedLists.cpp index 825ccfbb90..6ca6286145 100644 --- a/faiss/invlists/OnDiskInvertedLists.cpp +++ b/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& 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(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(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& 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(f); FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects"); FILE* fdesc = reader->f; diff --git a/faiss/invlists/OnDiskInvertedLists.h b/faiss/invlists/OnDiskInvertedLists.h index 98cb653a7a..40b0fff039 100644 --- a/faiss/invlists/OnDiskInvertedLists.h +++ b/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);