Skip to content

Commit

Permalink
Enable mmap-ped Faiss indices (IndexFlatCodes-based and HNSW) (zilliz…
Browse files Browse the repository at this point in the history 
…tech#996)

* HNSW_SQenable mmap-ped faiss indices (IndexFlatCodes-based and HNSW)

Signed-off-by: Alexandr Guzhva <[email protected]>

* temporary add IO_FLAG_MMAP_IFC

Signed-off-by: Alexandr Guzhva <[email protected]>

---------

Signed-off-by: Alexandr Guzhva <[email protected]>
  • Loading branch information
@alexanderguzhva
alexanderguzhva authored Dec 19, 2024
1 parent ca4ba32 commit c24fd21
Show file tree
Hide file tree
Showing 11 changed files with 507 additions and 27 deletions.
2 changes: 1 addition & 1 deletion src/index/flat/flat.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<IndexType, faiss::IndexFlat>::value) {
Expand Down
2 changes: 1 addition & 1 deletion src/index/hnsw/faiss_hnsw.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
3 changes: 2 additions & 1 deletion thirdparty/faiss/faiss/IndexFlatCodes.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,8 @@ CodePacker* IndexFlatCodes::get_CodePacker() const {
}

void IndexFlatCodes::permute_entries(const idx_t* perm) {
std::vector<uint8_t> new_codes(codes.size());
// std::vector<uint8_t> new_codes(codes.size());
MaybeOwnedVector<uint8_t> new_codes(codes.size());

for (idx_t i = 0; i < ntotal; i++) {
memcpy(new_codes.data() + i * code_size,
Expand Down
5 changes: 4 additions & 1 deletion thirdparty/faiss/faiss/IndexFlatCodes.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
#include <faiss/impl/DistanceComputer.h>
#include <vector>

#include <faiss/impl/maybe_owned_vector.h>

namespace faiss {

struct CodePacker;
Expand All @@ -23,7 +25,8 @@ struct IndexFlatCodes : Index {
size_t code_size;

/// encoded dataset, size ntotal * code_size
std::vector<uint8_t> codes;
// std::vector<uint8_t> codes;
MaybeOwnedVector<uint8_t> codes;

std::vector<float> code_norms;

Expand Down
3 changes: 2 additions & 1 deletion thirdparty/faiss/faiss/impl/HNSW.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

/**************************************************************
Expand Down
5 changes: 4 additions & 1 deletion thirdparty/faiss/faiss/impl/HNSW.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include <faiss/utils/Heap.h>
#include <faiss/utils/random.h>

#include <faiss/impl/maybe_owned_vector.h>

namespace faiss {

/** Implementation of the Hierarchical Navigable Small World
Expand Down Expand Up @@ -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<storage_idx_t> neighbors;
// std::vector<storage_idx_t> neighbors;
MaybeOwnedVector<storage_idx_t> neighbors;

/// entry point in the search structure (one of the points with maximum
/// level
Expand Down
151 changes: 131 additions & 20 deletions thirdparty/faiss/faiss/impl/index_read.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,90 @@
#include <faiss/IndexBinaryHash.h>
#include <faiss/IndexBinaryIVF.h>

// mmap facility
#include <faiss/impl/maybe_owned_vector.h>
#include <faiss/impl/mapped_io.h>

namespace faiss {

template<typename VectorT>
void read_vector(VectorT& target, IOReader* f) {
// is it a mmap-enabled reader?
MappedFileIOReader* mf = dynamic_cast<MappedFileIOReader*>(f);
if (mf != nullptr) {
// check if the use case is right
if constexpr(
std::is_same_v<VectorT, MaybeOwnedVector<uint8_t>> ||
std::is_same_v<VectorT, MaybeOwnedVector<float>> ||
std::is_same_v<VectorT, MaybeOwnedVector<int32_t>>
) {
// 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<typename VectorT>
void read_xb_vector(VectorT& target, IOReader* f) {
// is it a mmap-enabled reader?
MappedFileIOReader* mf = dynamic_cast<MappedFileIOReader*>(f);
if (mf != nullptr) {
// check if the use case is right
if constexpr(std::is_same_v<VectorT, MaybeOwnedVector<uint8_t>>) {
// 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
**************************************************************/
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand All @@ -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);
}
Expand Down Expand Up @@ -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(
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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);

Expand All @@ -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();
Expand Down Expand Up @@ -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
Expand All @@ -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")) {
Expand Down Expand Up @@ -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") ||
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<MmappedFileMappingOwner>(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<MmappedFileMappingOwner>(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
Expand Down
Loading

0 comments on commit c24fd21

Please sign in to comment.