Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Clean IVF_FLAT_NM compatible codes #856

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Clean IVF_FLAT_NM compatible codes #856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions include/knowhere/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,9 +180,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
Open in desktop
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,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 @@ -266,10 +266,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 @@ -1062,7 +1059,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 @@ -1083,53 +1080,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 @@ -1144,17 +1094,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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading