Skip to content

Commit

Permalink
Clean IVF_FLAT_NM compatible codes
Browse files Browse the repository at this point in the history 
Signed-off-by: Cai Yudong <[email protected]>
  • Loading branch information
@cydrain
cydrain committed Dec 19, 2024
1 parent ca4ba32 commit f1cebbf
Show file tree
Hide file tree
Showing 10 changed files with 4 additions and 310 deletions.
3 changes: 0 additions & 3 deletions include/knowhere/utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,9 +182,6 @@ round_down(const T value, const T align) {
return value / align * align;
}

extern void
ConvertIVFFlat(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, const size_t raw_size);

bool
UseDiskLoad(const std::string& index_type, const int32_t& /*version*/);

Expand Down
35 changes: 0 additions & 35 deletions src/common/utils.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,41 +87,6 @@ NormalizeDataset(const DataSetPtr dataset) {
NormalizeVecs<DataType>(data, rows, dim);
}

void
ConvertIVFFlat(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, const size_t raw_size) {
std::vector<std::string> names = {"IVF", // compatible with knowhere-1.x
knowhere::IndexEnum::INDEX_FAISS_IVFFLAT};
auto binary = binset.GetByNames(names);
if (binary == nullptr) {
return;
}

MemoryIOReader reader(binary->data.get(), binary->size);

try {
// only read IVF_FLAT index header
std::unique_ptr<faiss::IndexIVFFlat> ivfl;
ivfl.reset(static_cast<faiss::IndexIVFFlat*>(faiss::read_index_nm(&reader)));

// is_cosine is not defined in IVF_FLAT_NM, so mark it from config
ivfl->is_cosine = IsMetricType(metric_type, knowhere::metric::COSINE);

ivfl->restore_codes(raw_data, raw_size);

// over-write IVF_FLAT_NM binary with native IVF_FLAT binary
MemoryIOWriter writer;
faiss::write_index(ivfl.get(), &writer);
std::shared_ptr<uint8_t[]> data(writer.data());
binary->data = data;
binary->size = writer.tellg();

LOG_KNOWHERE_INFO_ << "Convert IVF_FLAT_NM to native IVF_FLAT, rows " << ivfl->ntotal << ", dim " << ivfl->d;
} catch (...) {
// not IVF_FLAT_NM format, do nothing
return;
}
}

bool
UseDiskLoad(const std::string& index_type, const int32_t& version) {
#ifdef KNOWHERE_WITH_CARDINAL
Expand Down
68 changes: 4 additions & 64 deletions src/index/ivf/ivf.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ class IvfIndexNode : public IndexNode {
}
Status
Serialize(BinarySet& binset) const override {
return this->SerializeImpl(binset, typename IndexDispatch<IndexType>::Tag{});
return this->SerializeImpl(binset);
}
Status
Deserialize(const BinarySet& binset, std::shared_ptr<Config> cfg) override;
Expand Down Expand Up @@ -294,10 +294,7 @@ class IvfIndexNode : public IndexNode {
GetIndexMetaImpl(std::unique_ptr<Config> cfg, IVFFlatTag) const;

Status
SerializeImpl(BinarySet& binset, IVFBaseTag) const;

Status
SerializeImpl(BinarySet& binset, IVFFlatTag) const;
SerializeImpl(BinarySet& binset) const;

Status
TrainInternal(const DataSetPtr dataset, std::shared_ptr<Config> cfg);
Expand Down Expand Up @@ -1089,7 +1086,7 @@ IvfIndexNode<DataType, IndexType>::GetIndexMetaImpl(std::unique_ptr<Config>, IVF

template <typename DataType, typename IndexType>
Status
IvfIndexNode<DataType, IndexType>::SerializeImpl(BinarySet& binset, IVFBaseTag) const {
IvfIndexNode<DataType, IndexType>::SerializeImpl(BinarySet& binset) const {
try {
if (!this->index_) {
LOG_KNOWHERE_WARNING_ << "index can not be serialized for empty index";
Expand All @@ -1110,53 +1107,6 @@ IvfIndexNode<DataType, IndexType>::SerializeImpl(BinarySet& binset, IVFBaseTag)
}
}

template <typename DataType, typename IndexType>
Status
IvfIndexNode<DataType, IndexType>::SerializeImpl(BinarySet& binset, IVFFlatTag) const {
try {
if (!this->index_) {
LOG_KNOWHERE_WARNING_ << "index can not be serialized for empty index";
return Status::empty_index;
}
MemoryIOWriter writer;
LOG_KNOWHERE_INFO_ << "request version " << this->version_.VersionNumber();
if (this->version_ <= Version::GetMinimalVersion()) {
faiss::write_index_nm(index_.get(), &writer);
LOG_KNOWHERE_INFO_ << "write IVF_FLAT_NM, file size " << writer.tellg();
} else {
faiss::write_index(index_.get(), &writer);
LOG_KNOWHERE_INFO_ << "write IVF_FLAT, file size " << writer.tellg();
}
std::shared_ptr<uint8_t[]> index_data_ptr(writer.data());
binset.Append(Type(), index_data_ptr, writer.tellg());

// append raw data for backward compatible
if (this->version_ <= Version::GetMinimalVersion()) {
size_t dim = index_->d;
size_t rows = index_->ntotal;
size_t raw_data_size = dim * rows * sizeof(float);
auto raw_data = std::make_unique<uint8_t[]>(raw_data_size);
for (size_t i = 0; i < index_->nlist; i++) {
size_t list_size = index_->invlists->list_size(i);
const faiss::idx_t* ids = index_->invlists->get_ids(i);
const uint8_t* codes = index_->invlists->get_codes(i);
for (size_t j = 0; j < list_size; j++) {
faiss::idx_t id = ids[j];
const uint8_t* src = codes + j * dim * sizeof(float);
uint8_t* dst = raw_data.get() + id * dim * sizeof(float);
memcpy(dst, src, dim * sizeof(float));
}
}
binset.Append("RAW_DATA", std::move(raw_data), raw_data_size);
LOG_KNOWHERE_INFO_ << "append raw data for IVF_FLAT_NM, size " << raw_data_size;
}
return Status::success;
} catch (const std::exception& e) {
LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what();
return Status::faiss_inner_error;
}
}

template <typename DataType, typename IndexType>
Status
IvfIndexNode<DataType, IndexType>::Deserialize(const BinarySet& binset, std::shared_ptr<Config> cfg) {
Expand All @@ -1171,17 +1121,7 @@ IvfIndexNode<DataType, IndexType>::Deserialize(const BinarySet& binset, std::sha

MemoryIOReader reader(binary->data.get(), binary->size);
try {
if constexpr (std::is_same<IndexType, faiss::IndexIVFFlat>::value) {
if (this->version_ <= Version::GetMinimalVersion()) {
auto raw_binary = binset.GetByName("RAW_DATA");
const BaseConfig& base_cfg = static_cast<const BaseConfig&>(*cfg);
ConvertIVFFlat(binset, base_cfg.metric_type.value(), raw_binary->data.get(), raw_binary->size);
// after conversion, binary size and data will be updated
reader.data_ = binary->data.get();
reader.total_ = binary->size;
}
index_.reset(static_cast<faiss::IndexIVFFlat*>(faiss::read_index(&reader)));
} else if constexpr (std::is_same<IndexType, faiss::IndexBinaryIVF>::value) {
if constexpr (std::is_same<IndexType, faiss::IndexBinaryIVF>::value) {
index_.reset(static_cast<IndexType*>(faiss::read_index_binary(&reader)));
} else {
index_.reset(static_cast<IndexType*>(faiss::read_index(&reader)));
Expand Down
7 changes: 0 additions & 7 deletions thirdparty/faiss/faiss/IndexIVFFlat.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,6 @@ IndexIVFFlat::IndexIVFFlat(
replace_invlists(new ArrayInvertedLists(nlist, code_size, is_cosine), true);
}

void IndexIVFFlat::restore_codes(
const uint8_t* raw_data,
const size_t raw_size) {
auto ails = dynamic_cast<faiss::ArrayInvertedLists*>(invlists);
ails->restore_codes(raw_data, raw_size, is_cosine);
}

void IndexIVFFlat::train(idx_t n, const float* x) {
if (is_cosine) {
auto x_normalized = knowhere::CopyAndNormalizeVecs(x, n, d);
Expand Down
2 changes: 0 additions & 2 deletions thirdparty/faiss/faiss/IndexIVFFlat.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,6 @@ struct IndexIVFFlat : IndexIVF {
MetricType = METRIC_L2,
bool is_cosine = false);

void restore_codes(const uint8_t* raw_data, const size_t raw_size);

// Be careful with overriding this function, because
// renormalized x may be used inside.
// Overridden by IndexIVFFlatDedup.
Expand Down
91 changes: 0 additions & 91 deletions thirdparty/faiss/faiss/impl/index_read.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -314,79 +314,6 @@ static void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
ivf->own_invlists = true;
}

InvertedLists *read_InvertedLists_nm(IOReader *f, int io_flags) {
uint32_t h;
READ1 (h);
if (h == fourcc("ilar") && !(io_flags & IO_FLAG_MMAP)) {
auto ails = new ArrayInvertedLists(0, 0);
READ1(ails->nlist);
READ1(ails->code_size);
ails->ids.resize(ails->nlist);
std::vector<size_t> sizes(ails->nlist);
read_ArrayInvertedLists_sizes(f, sizes);
for (size_t i = 0; i < ails->nlist; i++) {
ails->ids[i].resize(sizes[i]);
}
for (size_t i = 0; i < ails->nlist; i++) {
size_t n = ails->ids[i].size();
if (n > 0) {
READANDCHECK(ails->ids[i].data(), n);
}
}
return ails;
} else if (h == fourcc ("ilar") && (io_flags & IO_FLAG_MMAP)) {
// then we load it as an OnDiskInvertedLists
FileIOReader *reader = dynamic_cast<FileIOReader*>(f);
FAISS_THROW_IF_NOT_MSG(reader, "mmap only supported for File objects");
FILE *fdesc = reader->f;

auto ails = new OnDiskInvertedLists();
READ1(ails->nlist);
READ1(ails->code_size);
ails->read_only = true;
ails->lists.resize(ails->nlist);
std::vector<size_t> sizes(ails->nlist);
read_ArrayInvertedLists_sizes(f, sizes);
size_t o0 = ftell(fdesc), o = o0;
{ // do the mmap
struct stat buf;
int ret = fstat(fileno(fdesc), &buf);
FAISS_THROW_IF_NOT_FMT(ret == 0,
"fstat failed: %s", strerror(errno));
ails->totsize = buf.st_size;
ails->ptr = (uint8_t*)mmap(nullptr, ails->totsize,
PROT_READ, MAP_SHARED,
fileno(fdesc), 0);
FAISS_THROW_IF_NOT_FMT(ails->ptr != MAP_FAILED,
"could not mmap: %s",
strerror(errno));
madvise(ails->ptr, ails->totsize, MADV_RANDOM);
}

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);
}
FAISS_THROW_IF_NOT(o <= ails->totsize);
// resume normal reading of file
fseek (fdesc, o, SEEK_SET);
return ails;
} else {
FAISS_THROW_MSG("read_InvertedLists: unsupported invlist type");
}
}

void read_InvertedLists_nm(IndexIVF *ivf, IOReader *f, int io_flags) {
InvertedLists *ils = read_InvertedLists_nm (f, io_flags);
FAISS_THROW_IF_NOT(!ils || (ils->nlist == ivf->nlist &&
ils->code_size == ivf->code_size));
ivf->invlists = ils;
ivf->own_invlists = true;
}

static void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
READ1(pq->d);
READ1(pq->M);
Expand Down Expand Up @@ -1338,24 +1265,6 @@ Index* read_index(const char* fname, int io_flags) {
return idx;
}

// read offset-only index
Index *read_index_nm(IOReader *f, int io_flags) {
Index * idx = nullptr;
uint32_t h;
READ1(h);
if (h == fourcc("IwFl")) {
IndexIVFFlat * ivfl = new IndexIVFFlat ();
read_ivf_header (ivfl, f);
ivfl->code_size = ivfl->d * sizeof(float);
read_InvertedLists_nm(ivfl, f, io_flags);
idx = ivfl;
} else {
FAISS_THROW_FMT("Index type 0x%08x not supported\n", h);
idx = nullptr;
}
return idx;
}

VectorTransform* read_VectorTransform(const char* fname) {
FileIOReader reader(fname);
VectorTransform* vt = read_VectorTransform(&reader);
Expand Down
67 changes: 0 additions & 67 deletions thirdparty/faiss/faiss/impl/index_write.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -401,60 +401,6 @@ void write_InvertedLists(const InvertedLists* ils, IOWriter* f) {
}
}

// write inverted lists for offset-only index
void write_InvertedLists_nm(const InvertedLists *ils, IOWriter *f) {
if (ils == nullptr) {
uint32_t h = fourcc("il00");
WRITE1(h);
} else if (const auto & ails =
dynamic_cast<const ArrayInvertedLists *>(ils)) {
uint32_t h = fourcc("ilar");
WRITE1(h);
WRITE1(ails->nlist);
WRITE1(ails->code_size);
// here we store either as a full or a sparse data buffer
size_t n_non0 = 0;
for (size_t i = 0; i < ails->nlist; i++) {
if (ails->ids[i].size() > 0)
n_non0++;
}
if (n_non0 > ails->nlist / 2) {
uint32_t list_type = fourcc("full");
WRITE1(list_type);
std::vector<size_t> sizes;
for (size_t i = 0; i < ails->nlist; i++) {
sizes.push_back (ails->ids[i].size());
}
WRITEVECTOR(sizes);
} else {
int list_type = fourcc("sprs"); // sparse
WRITE1(list_type);
std::vector<size_t> sizes;
for (size_t i = 0; i < ails->nlist; i++) {
size_t n = ails->ids[i].size();
if (n > 0) {
sizes.push_back (i);
sizes.push_back (n);
}
}
WRITEVECTOR(sizes);
}
// make a single contiguous data buffer (useful for mmapping)
for (size_t i = 0; i < ails->nlist; i++) {
size_t n = ails->ids[i].size();
if (n > 0) {
// WRITEANDCHECK (ails->codes[i].data(), n * ails->code_size);
WRITEANDCHECK(ails->ids[i].data(), n);
}
}
} else {
fprintf(stderr, "WARN! write_InvertedLists: unsupported invlist type, "
"saving null invlist\n");
uint32_t h = fourcc("il00");
WRITE1(h);
}
}

void write_ProductQuantizer(const ProductQuantizer* pq, const char* fname) {
FileIOWriter writer(fname);
write_ProductQuantizer(pq, &writer);
Expand Down Expand Up @@ -1102,19 +1048,6 @@ void write_index(const Index* idx, const char* fname, int io_flags) {
write_index(idx, &writer, io_flags);
}

// write index for offset-only index
void write_index_nm(const Index *idx, IOWriter *f) {
if(const IndexIVFFlat * ivfl =
dynamic_cast<const IndexIVFFlat *> (idx)) {
uint32_t h = fourcc("IwFl");
WRITE1(h);
write_ivf_header(ivfl, f);
write_InvertedLists_nm(ivfl->invlists, f);
} else {
FAISS_THROW_MSG("don't know how to serialize this type of index");
}
}

void write_VectorTransform(const VectorTransform* vt, const char* fname) {
FileIOWriter writer(fname);
write_VectorTransform(vt, &writer);
Expand Down
3 changes: 0 additions & 3 deletions thirdparty/faiss/faiss/index_io.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,9 +90,6 @@ void write_ProductQuantizer(const ProductQuantizer* pq, IOWriter* f);
void write_InvertedLists(const InvertedLists* ils, IOWriter* f);
InvertedLists* read_InvertedLists(IOReader* reader, int io_flags = 0);

// for backward compatibility
Index *read_index_nm(IOReader *f, int io_flags = 0);
void write_index_nm(const Index* idx, IOWriter* writer);
} // namespace faiss

#endif
Loading

0 comments on commit f1cebbf

Please sign in to comment.