diff --git a/src/index/flat/flat.cc b/src/index/flat/flat.cc index 837352f2f..50075efcf 100644 --- a/src/index/flat/flat.cc +++ b/src/index/flat/flat.cc @@ -341,7 +341,7 @@ class FlatIndexNode : public IndexNode { int io_flags = 0; if (flat_cfg.enable_mmap.value()) { - io_flags |= faiss::IO_FLAG_MMAP; + io_flags |= faiss::IO_FLAG_MMAP_IFC; } if constexpr (std::is_same::value) { diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index c847e86ee..28f0307a8 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -213,7 +213,7 @@ class BaseFaissRegularIndexNode : public BaseFaissIndexNode { int io_flags = 0; if (cfg.enable_mmap.value()) { - io_flags |= faiss::IO_FLAG_MMAP; + io_flags |= faiss::IO_FLAG_MMAP_IFC; } try { diff --git a/thirdparty/faiss/faiss/IndexFlatCodes.cpp b/thirdparty/faiss/faiss/IndexFlatCodes.cpp index caff90ff9..897b350c0 100644 --- a/thirdparty/faiss/faiss/IndexFlatCodes.cpp +++ b/thirdparty/faiss/faiss/IndexFlatCodes.cpp @@ -104,7 +104,8 @@ CodePacker* IndexFlatCodes::get_CodePacker() const { } void IndexFlatCodes::permute_entries(const idx_t* perm) { - std::vector new_codes(codes.size()); + // std::vector new_codes(codes.size()); + MaybeOwnedVector new_codes(codes.size()); for (idx_t i = 0; i < ntotal; i++) { memcpy(new_codes.data() + i * code_size, diff --git a/thirdparty/faiss/faiss/IndexFlatCodes.h b/thirdparty/faiss/faiss/IndexFlatCodes.h index ad91c8601..414c863bf 100644 --- a/thirdparty/faiss/faiss/IndexFlatCodes.h +++ b/thirdparty/faiss/faiss/IndexFlatCodes.h @@ -13,6 +13,8 @@ #include #include +#include + namespace faiss { struct CodePacker; @@ -23,7 +25,8 @@ struct IndexFlatCodes : Index { size_t code_size; /// encoded dataset, size ntotal * code_size - std::vector codes; + // std::vector codes; + MaybeOwnedVector codes; std::vector code_norms; diff --git a/thirdparty/faiss/faiss/impl/HNSW.cpp b/thirdparty/faiss/faiss/impl/HNSW.cpp index 4e6fb2944..5deeeacb8 100644 --- a/thirdparty/faiss/faiss/impl/HNSW.cpp +++ b/thirdparty/faiss/faiss/impl/HNSW.cpp @@ -1156,7 +1156,8 @@ void HNSW::permute_entries(const idx_t* map) { // swap everyone std::swap(levels, new_levels); std::swap(offsets, new_offsets); - std::swap(neighbors, new_neighbors); + // std::swap(neighbors, new_neighbors); + neighbors = std::move(new_neighbors); } /************************************************************** diff --git a/thirdparty/faiss/faiss/impl/HNSW.h b/thirdparty/faiss/faiss/impl/HNSW.h index f376c9fcc..eea2ed627 100644 --- a/thirdparty/faiss/faiss/impl/HNSW.h +++ b/thirdparty/faiss/faiss/impl/HNSW.h @@ -19,6 +19,8 @@ #include #include +#include + namespace faiss { /** Implementation of the Hierarchical Navigable Small World @@ -120,7 +122,8 @@ struct HNSW { /// neighbors[offsets[i]:offsets[i+1]] is the list of neighbors of vector i /// for all levels. this is where all storage goes. - std::vector neighbors; + // std::vector neighbors; + MaybeOwnedVector neighbors; /// entry point in the search structure (one of the points with maximum /// level diff --git a/thirdparty/faiss/faiss/impl/index_read.cpp b/thirdparty/faiss/faiss/impl/index_read.cpp index 599953208..c2bf84b0c 100644 --- a/thirdparty/faiss/faiss/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/impl/index_read.cpp @@ -59,8 +59,90 @@ #include #include +// mmap facility +#include +#include + namespace faiss { +template +void read_vector(VectorT& target, IOReader* f) { + // is it a mmap-enabled reader? + MappedFileIOReader* mf = dynamic_cast(f); + if (mf != nullptr) { + // check if the use case is right + if constexpr( + std::is_same_v> || + std::is_same_v> || + std::is_same_v> + ) { + // read the size + size_t size = 0; + READANDCHECK(&size, 1); + // ok, mmap and check + char* address = nullptr; + const size_t nread = mf->mmap((void**)&address, sizeof(typename VectorT::value_type), size); + + FAISS_THROW_IF_NOT_FMT( + nread == (size), + "read error in %s: %zd != %zd (%s)", + f->name.c_str(), + nread, + size, + strerror(errno)); + + VectorT mmapped = VectorT::from_mmapped( + address, nread, mf->mmap_owner + ); + target = std::move(mmapped); + + return; + } + } + + // the default case + READVECTOR(target); +} + +template +void read_xb_vector(VectorT& target, IOReader* f) { + // is it a mmap-enabled reader? + MappedFileIOReader* mf = dynamic_cast(f); + if (mf != nullptr) { + // check if the use case is right + if constexpr(std::is_same_v>) { + // read the size + size_t size = 0; + READANDCHECK(&size, 1); + + size *= 4; + + // ok, mmap and check + char* address = nullptr; + const size_t nread = mf->mmap((void**)&address, sizeof(typename VectorT::value_type), size); + + FAISS_THROW_IF_NOT_FMT( + nread == (size), + "read error in %s: %zd != %zd (%s)", + f->name.c_str(), + nread, + size, + strerror(errno)); + + VectorT mmapped = VectorT::from_mmapped( + address, nread, mf->mmap_owner + ); + target = std::move(mmapped); + + return; + } + } + + // the default case + READXBVECTOR(target); +} + + /************************************************************* * Read **************************************************************/ @@ -512,7 +594,8 @@ static void read_HNSW(HNSW* hnsw, IOReader* f) { READVECTOR(hnsw->cum_nneighbor_per_level); READVECTOR(hnsw->levels); READVECTOR(hnsw->offsets); - READVECTOR(hnsw->neighbors); + // READVECTOR(hnsw->neighbors); + read_vector(hnsw->neighbors, f); READ1(hnsw->entry_point); READ1(hnsw->max_level); @@ -690,7 +773,8 @@ Index* read_index(IOReader* f, int io_flags) { IndexFlatCosine* idxf = new IndexFlatCosine(); read_index_header(idxf, f); idxf->code_size = idxf->d * sizeof(float); - READXBVECTOR(idxf->codes); + // READXBVECTOR(idxf->codes); + read_xb_vector(idxf->codes, f); READVECTOR(idxf->code_norms); // reconstruct inverse norms @@ -712,7 +796,8 @@ Index* read_index(IOReader* f, int io_flags) { } read_index_header(idxf, f); idxf->code_size = idxf->d * sizeof(float); - READXBVECTOR(idxf->codes); + // READXBVECTOR(idxf->codes); + read_xb_vector(idxf->codes, f); if (idxf->is_cosine) { READVECTOR(idxf->code_norms); } @@ -746,7 +831,8 @@ Index* read_index(IOReader* f, int io_flags) { idxl->rrot = *rrot; delete rrot; } - READVECTOR(idxl->codes); + // READVECTOR(idxl->codes); + read_vector(idxl->codes, f); FAISS_THROW_IF_NOT( idxl->rrot.d_in == idxl->d && idxl->rrot.d_out == idxl->nbits); FAISS_THROW_IF_NOT( @@ -757,7 +843,8 @@ Index* read_index(IOReader* f, int io_flags) { read_index_header(idxp, f); read_ProductQuantizer(&idxp->pq, f); idxp->code_size = idxp->pq.code_size; - READVECTOR(idxp->codes); + // READVECTOR(idxp->codes); + read_vector(idxp->codes, f); READ1(idxp->search_type); READ1(idxp->encode_signs); READ1(idxp->polysemous_ht); @@ -776,7 +863,8 @@ Index* read_index(IOReader* f, int io_flags) { read_index_header(idxp, f); read_ProductQuantizer(&idxp->pq, f); idxp->code_size = idxp->pq.code_size; - READVECTOR(idxp->codes); + // READVECTOR(idxp->codes); + read_vector(idxp->codes, f); if (h == fourcc("IxPo") || h == fourcc("IxPq")) { READ1(idxp->search_type); READ1(idxp->encode_signs); @@ -804,21 +892,24 @@ Index* read_index(IOReader* f, int io_flags) { read_ResidualQuantizer(&idxr->rq, f, io_flags); } READ1(idxr->code_size); - READVECTOR(idxr->codes); + // READVECTOR(idxr->codes); + read_vector(idxr->codes, f); idx = idxr; } else if (h == fourcc("IxLS")) { auto idxr = new IndexLocalSearchQuantizer(); read_index_header(idxr, f); read_LocalSearchQuantizer(&idxr->lsq, f); READ1(idxr->code_size); - READVECTOR(idxr->codes); + // READVECTOR(idxr->codes); + read_vector(idxr->codes, f); idx = idxr; } else if (h == fourcc("IxP5")) { auto idxpr = new IndexProductResidualQuantizerCosine(); read_index_header(idxpr, f); read_ProductResidualQuantizer(&idxpr->prq, f, io_flags); READ1(idxpr->code_size); - READVECTOR(idxpr->codes); + // READVECTOR(idxpr->codes); + read_vector(idxpr->codes, f); // read inverse norms READVECTOR(idxpr->inverse_norms_storage.inverse_l2_norms); @@ -828,14 +919,16 @@ Index* read_index(IOReader* f, int io_flags) { read_index_header(idxpr, f); read_ProductResidualQuantizer(&idxpr->prq, f, io_flags); READ1(idxpr->code_size); - READVECTOR(idxpr->codes); + // READVECTOR(idxpr->codes); + read_vector(idxpr->codes, f); idx = idxpr; } else if (h == fourcc("IxPL")) { auto idxpl = new IndexProductLocalSearchQuantizer(); read_index_header(idxpl, f); read_ProductLocalSearchQuantizer(&idxpl->plsq, f); READ1(idxpl->code_size); - READVECTOR(idxpl->codes); + // READVECTOR(idxpl->codes); + read_vector(idxpl->codes, f); idx = idxpl; } else if (h == fourcc("ImRQ")) { ResidualCoarseQuantizer* idxr = new ResidualCoarseQuantizer(); @@ -1004,7 +1097,8 @@ Index* read_index(IOReader* f, int io_flags) { IndexScalarQuantizerCosine* idxs = new IndexScalarQuantizerCosine(); read_index_header(idxs, f); read_ScalarQuantizer(&idxs->sq, f); - READVECTOR(idxs->codes); + // READVECTOR(idxs->codes); + read_vector(idxs->codes, f); idxs->code_size = idxs->sq.code_size; // reconstruct inverse norms @@ -1015,7 +1109,8 @@ Index* read_index(IOReader* f, int io_flags) { IndexScalarQuantizer* idxs = new IndexScalarQuantizer(); read_index_header(idxs, f); read_ScalarQuantizer(&idxs->sq, f); - READVECTOR(idxs->codes); + // READVECTOR(idxs->codes); + read_vector(idxs->codes, f); idxs->code_size = idxs->sq.code_size; idx = idxs; } else if (h == fourcc("IxLa")) { @@ -1188,7 +1283,8 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxp->code_size_1); READ1(idxp->code_size_2); READ1(idxp->code_size); - READVECTOR(idxp->codes); + // READVECTOR(idxp->codes); + read_vector(idxp->codes, f); idx = idxp; } else if ( h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") || @@ -1266,7 +1362,8 @@ Index* read_index(IOReader* f, int io_flags) { READ1(idxpqfs->qbs); READ1(idxpqfs->ntotal2); READ1(idxpqfs->M2); - READVECTOR(idxpqfs->codes); + // READVECTOR(idxpqfs->codes); + read_vector(idxpqfs->codes, f); const auto& pq = idxpqfs->pq; idxpqfs->M = pq.M; @@ -1328,14 +1425,28 @@ Index* read_index(IOReader* f, int io_flags) { } Index* read_index(FILE* f, int io_flags) { - FileIOReader reader(f); - return read_index(&reader, io_flags); + if ((io_flags & IO_FLAG_MMAP_IFC) == IO_FLAG_MMAP_IFC) { + // enable mmap-supporting IOReader + auto owner = std::make_shared(f); + MappedFileIOReader reader(owner); + return read_index(&reader, io_flags); + } else { + FileIOReader reader(f); + return read_index(&reader, io_flags); + } } Index* read_index(const char* fname, int io_flags) { - FileIOReader reader(fname); - Index* idx = read_index(&reader, io_flags); - return idx; + if ((io_flags & IO_FLAG_MMAP_IFC) == IO_FLAG_MMAP_IFC) { + // enable mmap-supporting IOReader + auto owner = std::make_shared(fname); + MappedFileIOReader reader(owner); + return read_index(&reader, io_flags); + } else { + FileIOReader reader(fname); + Index* idx = read_index(&reader, io_flags); + return idx; + } } // read offset-only index diff --git a/thirdparty/faiss/faiss/impl/mapped_io.cpp b/thirdparty/faiss/faiss/impl/mapped_io.cpp new file mode 100644 index 000000000..6e79031e7 --- /dev/null +++ b/thirdparty/faiss/faiss/impl/mapped_io.cpp @@ -0,0 +1,127 @@ +#include + +// this is the linux version +// todo: add Windows implementation +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include + +namespace faiss { + +struct MmappedFileMappingOwner::PImpl { + void* ptr = nullptr; + size_t ptr_size = 0; + + PImpl(FILE* f) { + // get the size + struct stat s; + int status = fstat(fileno(f), &s); + FAISS_THROW_IF_NOT_FMT( + status >= 0, "fstat() failed: %s", strerror(errno) + ); + + const size_t filesize = s.st_size; + + void* address = mmap(nullptr, filesize, PROT_READ, MAP_SHARED, fileno(f), 0); + FAISS_THROW_IF_NOT_FMT( + address != nullptr, "could not mmap(): %s", strerror(errno) + ); + + // btw, fd can be closed here + + // set 'random' access pattern + // todo: check the error + madvise(address, filesize, MADV_RANDOM); + + // save it + ptr = address; + ptr_size = filesize; + } + + ~PImpl() { + // todo: check for an error + munmap(ptr, ptr_size); + } +}; + +MmappedFileMappingOwner::MmappedFileMappingOwner(const std::string& filename) { + auto fd = std::unique_ptr(fopen(filename.c_str(), "r"), &fclose); + FAISS_THROW_IF_NOT_FMT( + fd.get(), "could not open %s for reading: %s", filename.c_str(), strerror(errno)); + + p_impl = std::make_unique(fd.get()); +} + +MmappedFileMappingOwner::MmappedFileMappingOwner(FILE* f) { + p_impl = std::make_unique(f); +} + +MmappedFileMappingOwner::~MmappedFileMappingOwner() = default; + +// +void* MmappedFileMappingOwner::data() const { + return p_impl->ptr; +} + +size_t MmappedFileMappingOwner::size() const { + return p_impl->ptr_size; +} + + + + +MappedFileIOReader::MappedFileIOReader(const std::shared_ptr& owner) : + mmap_owner(owner) {} + +// this operation performs a copy +size_t MappedFileIOReader::operator()(void* ptr, size_t size, size_t nitems) { + char* ptr_c = nullptr; + + const size_t actual_nitems = this->mmap((void**)&ptr_c, size, nitems); + if (actual_nitems > 0) { + memcpy(ptr, ptr_c, size * actual_nitems); + } + + return actual_nitems; +} + +// this operation returns a mmapped address, owned by mmap_owner +size_t MappedFileIOReader::mmap(void** ptr, size_t size, size_t nitems) { + if (size == 0) { + return nitems; + } + + size_t actual_size = size * nitems; + if (pos + size * nitems > mmap_owner->size()) { + actual_size = mmap_owner->size() - pos; + } + + size_t actual_nitems = (actual_size + size - 1) / size; + if (actual_nitems == 0) { + return 0; + } + + // get an address + *ptr = (void*)(reinterpret_cast(mmap_owner->data()) + pos); + + // alter pos + pos += size * actual_nitems; + + return actual_nitems; +} + +int MappedFileIOReader::filedescriptor() { + // todo + return -1; +} + +} \ No newline at end of file diff --git a/thirdparty/faiss/faiss/impl/mapped_io.h b/thirdparty/faiss/faiss/impl/mapped_io.h new file mode 100644 index 000000000..d797b65c7 --- /dev/null +++ b/thirdparty/faiss/faiss/impl/mapped_io.h @@ -0,0 +1,42 @@ +#pragma once + +#include +#include +#include + +#include +#include + +namespace faiss { + +// holds a memory-mapped region over a file +struct MmappedFileMappingOwner : public MappingOwner { + MmappedFileMappingOwner(const std::string& filename); + MmappedFileMappingOwner(FILE* f); + ~MmappedFileMappingOwner(); + + void* data() const; + size_t size() const; + + struct PImpl; + std::unique_ptr p_impl; +}; + +// a deserializer that supports memory-mapped files +struct MappedFileIOReader : IOReader { + std::shared_ptr mmap_owner; + + size_t pos = 0; + + MappedFileIOReader(const std::shared_ptr& owner); + + // perform a copy + size_t operator()(void* ptr, size_t size, size_t nitems) override; + // perform a quasi-read that returns a mmapped address, owned by mmap_owner, + // and updates the position + size_t mmap(void** ptr, size_t size, size_t nitems); + + int filedescriptor() override; +}; + +} \ No newline at end of file diff --git a/thirdparty/faiss/faiss/impl/maybe_owned_vector.h b/thirdparty/faiss/faiss/impl/maybe_owned_vector.h new file mode 100644 index 000000000..0437a6195 --- /dev/null +++ b/thirdparty/faiss/faiss/impl/maybe_owned_vector.h @@ -0,0 +1,188 @@ +#pragma once + +#include +#include +#include +#include + +#include + +namespace faiss { + +struct MappingOwner { + virtual ~MappingOwner() = default; +}; + +// a container that either works as std::vector that owns its own memory, +// or as a mapped pointer owned by someone third-party owner +template +struct MaybeOwnedVector { + using value_type = T; + using self_type = MaybeOwnedVector; + + bool is_owned = true; + + // this one is used if is_owned == true + std::vector owned_data; + + // these three are used if is_owned == false + T* mapped_data = nullptr; + // the number of T elements + size_t mapped_size = 0; + std::shared_ptr mapping_owner; + + // points either to mapped_data, or to owned.data() + T* c_ptr = nullptr; + // uses either mapped_size, or owned.size(); + size_t c_size = 0; + + MaybeOwnedVector() = default; + MaybeOwnedVector(const size_t initial_size) { + is_owned = true; + + owned_data.resize(initial_size); + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } + + MaybeOwnedVector(const MaybeOwnedVector& other) { + is_owned = other.is_owned; + owned_data = other.owned_data; + + mapped_data = other.mapped_data; + mapped_size = other.mapped_size; + mapping_owner = other.mapping_owner; + + if (is_owned) { + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } else { + c_ptr = mapped_data; + c_size = mapped_size; + } + } + + MaybeOwnedVector(MaybeOwnedVector&& other) { + is_owned = other.is_owned; + owned_data = std::move(other.owned_data); + + mapped_data = other.mapped_data; + mapped_size = other.mapped_size; + mapping_owner = std::move(other.mapping_owner); + + if (is_owned) { + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } else { + c_ptr = mapped_data; + c_size = mapped_size; + } + } + + MaybeOwnedVector& operator =(const MaybeOwnedVector& other) { + if (this == &other) { + return *this; + } + + // create a copy + MaybeOwnedVector cloned(other); + // swap + swap(*this, cloned); + + return *this; + } + + MaybeOwnedVector& operator =(MaybeOwnedVector&& other) { + if (this == &other) { + return *this; + } + + // moved + MaybeOwnedVector moved(std::move(other)); + // swap + swap(*this, moved); + + return *this; + } + + MaybeOwnedVector(std::vector&& other) { + is_owned = true; + + owned_data = std::move(other); + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } + + static MaybeOwnedVector from_mmapped( + void* address, + const size_t n_mapped_elements, + const std::shared_ptr& owner + ) { + MaybeOwnedVector vec; + vec.is_owned = false; + vec.mapped_data = reinterpret_cast(address); + vec.mapped_size = n_mapped_elements; + vec.mapping_owner = owner; + + vec.c_ptr = vec.mapped_data; + vec.c_size = vec.mapped_size; + + return vec; + } + + const T* data() const { + return c_ptr; + } + + T* data() { + return c_ptr; + } + + size_t size() const { + return c_size; + } + + T& operator[](const size_t idx) { + return c_ptr[idx]; + } + + const T& operator[](const size_t idx) const { + return c_ptr[idx]; + } + + void clear() { + FAISS_ASSERT_MSG(is_owned, "This operation cannot be performed on a memory-mapped vector"); + + owned_data.clear(); + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } + + void resize(const size_t new_size) { + FAISS_ASSERT_MSG(is_owned, "This operation cannot be performed on a memory-mapped vector"); + + owned_data.resize(new_size); + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } + + void resize(const size_t new_size, const value_type v) { + FAISS_ASSERT_MSG(is_owned, "This operation cannot be performed on a memory-mapped vector"); + + owned_data.resize(new_size, v); + c_ptr = owned_data.data(); + c_size = owned_data.size(); + } + + friend void swap(self_type& a, self_type& b) { + std::swap(a.is_owned, b.is_owned); + std::swap(a.owned_data, b.owned_data); + std::swap(a.mapped_data, b.mapped_data); + std::swap(a.mapped_size, b.mapped_size); + std::swap(a.mapping_owner, b.mapping_owner); + std::swap(a.c_ptr, b.c_ptr); + std::swap(a.c_size, b.c_size); + } +}; + +} \ No newline at end of file diff --git a/thirdparty/faiss/faiss/index_io.h b/thirdparty/faiss/faiss/index_io.h index a78b1493f..8e08746ca 100644 --- a/thirdparty/faiss/faiss/index_io.h +++ b/thirdparty/faiss/faiss/index_io.h @@ -64,8 +64,12 @@ const int IO_FLAG_PQ_SKIP_SDC_TABLE = 32; // load index data with vectors' norms const int IO_FLAG_WITH_NORM = 1 << 8; // try to memmap data (useful to load an ArrayInvertedLists as an -// OnDiskInvertedLists) +// OnDiskInvertedLists). const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000; +// mmap that handles codes for IndexFlatCodes-derived indices and HNSW. +// this is a temporary solution, it is expected to be merged with IO_FLAG_MMAP +// after OnDiskInvertedLists get properly updated. +const int IO_FLAG_MMAP_IFC = 1 << 9; Index* read_index(const char* fname, int io_flags = 0); Index* read_index(FILE* f, int io_flags = 0);