milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Optimize functions (#1803) 

Signed-off-by: sahuang <xiaohai.xu@zilliz.com>
pull/1811/head
Xiaohai Xu 2020-03-30 18:11:08 +08:00 committed by GitHub
parent dbbf91db28
commit 88cf5e2bd8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 12 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.cu vendored
View File

@ -212,13 +212,11 @@ GpuIndexIVFFlat::addImpl_(int n,
// Data is already resident on the GPU
Tensor<float, 2, true> data(const_cast<float*>(x), {n, (int) this->d});
auto bitset = toDevice<uint8_t, 1>(resources_, device_, nullptr, stream, {0});
static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
Tensor<long, 1, true> labels(const_cast<long*>(xids), {n});
// Not all vectors may be able to be added (some may contain NaNs etc)
index_->classifyAndAddVectors(data, labels, bitset);
index_->classifyAndAddVectors(data, labels);
// but keep the ntotal based on the total number of vectors that we attempted
// to add

4
core/src/index/thirdparty/faiss/gpu/GpuIndexIVFPQ.cu vendored
View File

@ -335,13 +335,11 @@ GpuIndexIVFPQ::addImpl_(int n,
// Data is already resident on the GPU
Tensor<float, 2, true> data(const_cast<float*>(x), {n, (int) this->d});
auto bitset = toDevice<uint8_t, 1>(resources_, device_, nullptr, stream, {0});
static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
Tensor<long, 1, true> labels(const_cast<long*>(xids), {n});
// Not all vectors may be able to be added (some may contain NaNs etc)
index_->classifyAndAddVectors(data, labels, bitset);
index_->classifyAndAddVectors(data, labels);
// but keep the ntotal based on the total number of vectors that we attempted
// to add

4
core/src/index/thirdparty/faiss/gpu/GpuIndexIVFSQHybrid.cu vendored
View File

@ -309,13 +309,11 @@ GpuIndexIVFSQHybrid::addImpl_(int n,
// Data is already resident on the GPU
Tensor<float, 2, true> data(const_cast<float*>(x), {n, (int) this->d});
auto bitset = toDevice<uint8_t, 1>(resources_, device_, nullptr, stream, {0});
static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
Tensor<long, 1, true> labels(const_cast<long*>(xids), {n});
// Not all vectors may be able to be added (some may contain NaNs etc)
index_->classifyAndAddVectors(data, labels, bitset);
index_->classifyAndAddVectors(data, labels);
// but keep the ntotal based on the total number of vectors that we attempted
// to add

4
core/src/index/thirdparty/faiss/gpu/GpuIndexIVFScalarQuantizer.cu vendored
View File

@ -244,13 +244,11 @@ GpuIndexIVFScalarQuantizer::addImpl_(int n,
// Data is already resident on the GPU
Tensor<float, 2, true> data(const_cast<float*>(x), {n, (int) this->d});
auto bitset = toDevice<uint8_t, 1>(resources_, device_, nullptr, stream, {0});
static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
Tensor<long, 1, true> labels(const_cast<long*>(xids), {n});
// Not all vectors may be able to be added (some may contain NaNs etc)
index_->classifyAndAddVectors(data, labels, bitset);
index_->classifyAndAddVectors(data, labels);
// but keep the ntotal based on the total number of vectors that we attempted
// to add

5
core/src/index/thirdparty/faiss/gpu/impl/IVFFlat.cu vendored
View File

@ -157,14 +157,15 @@ IVFFlat::addCodeVectorsFromCpu(int listId,
int
IVFFlat::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
Tensor<long, 1, true>& indices,
Tensor<uint8_t, 1, true>& bitset) {
Tensor<long, 1, true>& indices) {
FAISS_ASSERT(vecs.getSize(0) == indices.getSize(0));
FAISS_ASSERT(vecs.getSize(1) == dim_);
auto& mem = resources_->getMemoryManagerCurrentDevice();
auto stream = resources_->getDefaultStreamCurrentDevice();
DeviceTensor<uint8_t, 1, true> bitset(mem, {0}, stream);
// Number of valid vectors that we actually add; we return this
int numAdded = 0;

3
core/src/index/thirdparty/faiss/gpu/impl/IVFFlat.cuh vendored
View File

@ -44,8 +44,7 @@ class IVFFlat : public IVFBase {
/// Returns the number of vectors successfully added. Vectors may
/// not be able to be added because they contain NaNs.
int classifyAndAddVectors(Tensor<float, 2, true>& vecs,
Tensor<long, 1, true>& indices,
Tensor<uint8_t, 1, true>& bitset);
Tensor<long, 1, true>& indices);
/// Find the approximate k nearest neigbors for `queries` against

5
core/src/index/thirdparty/faiss/gpu/impl/IVFPQ.cu vendored
View File

@ -110,8 +110,7 @@ IVFPQ::setPrecomputedCodes(bool enable) {
int
IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
Tensor<long, 1, true>& indices,
Tensor<uint8_t, 1, true>& bitset) {
Tensor<long, 1, true>& indices) {
FAISS_ASSERT(vecs.getSize(0) == indices.getSize(0));
FAISS_ASSERT(vecs.getSize(1) == dim_);
@ -120,6 +119,8 @@ IVFPQ::classifyAndAddVectors(Tensor<float, 2, true>& vecs,
auto& mem = resources_->getMemoryManagerCurrentDevice();
auto stream = resources_->getDefaultStreamCurrentDevice();
DeviceTensor<uint8_t, 1, true> bitset(mem, {0}, stream);
// Number of valid vectors that we actually add; we return this
int numAdded = 0;

3
core/src/index/thirdparty/faiss/gpu/impl/IVFPQ.cuh vendored
View File

@ -52,8 +52,7 @@ class IVFPQ : public IVFBase {
/// Returns the number of vectors successfully added. Vectors may
/// not be able to be added because they contain NaNs.
int classifyAndAddVectors(Tensor<float, 2, true>& vecs,
Tensor<long, 1, true>& indices,
Tensor<uint8_t, 1, true>& bitset);
Tensor<long, 1, true>& indices);
/// Find the approximate k nearest neigbors for `queries` against
/// our database