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

fix search by GPU (#2453) 

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
pull/2536/head
shengjun.li 2020-05-28 23:39:16 +08:00 committed by JinHai-CN
parent 87d22da154
commit 96c68edcb6
5 changed files with 11 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@ -34,6 +34,7 @@ Please mark all change in change log and use the issue from GitHub
- \#2395 Fix large nq cudaMalloc error
- \#2399 The nlist set by the user may not take effect
- \#2403 MySQL max_idle_time is 10 by default
- \#2450 The deleted vectors may be found on GPU
## Feature

7
core/src/index/thirdparty/faiss/gpu/impl/IVFUtilsSelect1.cu vendored
View File

@ -75,10 +75,9 @@ pass1SelectLists(void** listIndices,
topQueryToCentroid,
opt);
if (bitsetEmpty || (!(bitset[index >> 3] & (0x1 << (index & 0x7))))) {
heap.add(distanceStart[i], start + i);
} else {
heap.add((1.0 / 0.0), start + i);
heap.addThreadQ(distanceStart[i], start + i);
}
heap.checkThreadQ();
}
// Handle warp divergence separately
@ -91,8 +90,6 @@ pass1SelectLists(void** listIndices,
opt);
if (bitsetEmpty || (!(bitset[index >> 3] & (0x1 << (index & 0x7))))) {
heap.addThreadQ(distanceStart[i], start + i);
} else {
heap.addThreadQ((1.0 / 0.0), start + i);
}
}

9
core/src/index/thirdparty/faiss/gpu/impl/L2Select.cu vendored
View File

@ -161,21 +161,18 @@ __global__ void l2SelectMinK(Tensor<T, 2, true> productDistances,
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
v = Math<T>::add(centroidDistances[i],
productDistances[row][i]);
} else {
v = (T)(1.0 / 0.0);
heap.addThreadQ(v, i);
}
heap.add(v, i);
heap.checkThreadQ();
}
if (i < productDistances.getSize(1)) {
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
v = Math<T>::add(centroidDistances[i],
productDistances[row][i]);
} else {
v = (T)(1.0 / 0.0);
}
heap.addThreadQ(v, i);
}
}
heap.reduce();
for (int i = threadIdx.x; i < k; i += blockDim.x) {

14
core/src/index/thirdparty/faiss/gpu/utils/BlockSelectKernel.cuh vendored
View File

@ -146,10 +146,9 @@ __global__ void blockSelect(Tensor<K, 2, true> in,
for (; i < limit; i += ThreadsPerBlock) {
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
heap.add(*inStart, (IndexType) i);
} else {
heap.add(-1.0, (IndexType) i);
heap.addThreadQ(*inStart, (IndexType) i);
}
heap.checkThreadQ();
inStart += ThreadsPerBlock;
}
@ -158,8 +157,6 @@ __global__ void blockSelect(Tensor<K, 2, true> in,
if (i < in.getSize(1)) {
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
heap.addThreadQ(*inStart, (IndexType) i);
} else {
heap.addThreadQ(-1.0, (IndexType) i);
}
}
@ -208,10 +205,9 @@ __global__ void blockSelectPair(Tensor<K, 2, true> inK,
for (; i < limit; i += ThreadsPerBlock) {
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
heap.add(*inKStart, *inVStart);
} else {
heap.add(-1.0, *inVStart);
heap.addThreadQ(*inKStart, *inVStart);
}
heap.checkThreadQ();
inKStart += ThreadsPerBlock;
inVStart += ThreadsPerBlock;
@ -221,8 +217,6 @@ __global__ void blockSelectPair(Tensor<K, 2, true> inK,
if (i < inK.getSize(1)) {
if (bitsetEmpty || (!(bitset[i >> 3] & (0x1 << (i & 0x7))))) {
heap.addThreadQ(*inKStart, *inVStart);
} else {
heap.addThreadQ(-1.0, *inVStart);
}
}

9
core/src/scheduler/task/SearchTask.cpp
View File

@ -288,15 +288,6 @@ XSearchTask::Execute() {
{
std::unique_lock<std::mutex> lock(search_job->mutex());
if (search_job->GetResultIds().size() > spec_k) {
if (search_job->GetResultIds().front() == -1) {
// initialized results set
search_job->GetResultIds().resize(spec_k * nq);
search_job->GetResultDistances().resize(spec_k * nq);
}
}
search_job->vector_count() = nq;
XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, ascending_reduce,
search_job->GetResultIds(), search_job->GetResultDistances());