diff --git a/c_api/index_io_c_ex.h b/c_api/index_io_c_ex.h index 77a3f999b1..7c85409b8c 100644 --- a/c_api/index_io_c_ex.h +++ b/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); 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..cf967e8034 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,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(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(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")) { diff --git a/faiss/index_io.h b/faiss/index_io.h index 8d52ee1afd..0fd898edd6 100644 --- a/faiss/index_io.h +++ b/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; diff --git a/faiss/invlists/OnDiskInvertedLists.cpp b/faiss/invlists/OnDiskInvertedLists.cpp index 825ccfbb90..2896190447 100644 --- a/faiss/invlists/OnDiskInvertedLists.cpp +++ b/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& 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(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(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& 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(f); FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects"); diff --git a/faiss/invlists/OnDiskInvertedLists.h b/faiss/invlists/OnDiskInvertedLists.h index 98cb653a7a..65f07b5c40 100644 --- a/faiss/invlists/OnDiskInvertedLists.h +++ b/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);