Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
BaiXuePrincess
milvus
提交
6dd89bb0
milvus
项目概览
BaiXuePrincess
/
milvus
与 Fork 源项目一致
从无法访问的项目Fork
通知
7
Star
4
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
6dd89bb0
编写于
9月 09, 2019
作者:
X
xj.lin
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
MS-534 1. fix
Former-commit-id: 6d1927610d68cdb2689070699e42851223ee67d9
上级
724be3da
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
71 addition
and
120 deletion
+71
-120
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
+18
-16
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
+47
-98
cpp/src/core/test/test_ivf.cpp
cpp/src/core/test/test_ivf.cpp
+1
-1
cpp/src/utils/BlockingQueue.inl
cpp/src/utils/BlockingQueue.inl
+5
-5
未找到文件。
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
浏览文件 @
6dd89bb0
...
...
@@ -3,6 +3,7 @@
#include <faiss/gpu/StandardGpuResources.h>
#include "ivf.h"
#include "src/utils/BlockingQueue.h"
namespace
zilliz
{
...
...
@@ -16,12 +17,15 @@ struct Resource {
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
faiss_res
;
int64_t
id
;
std
::
mutex
mutex
;
};
using
ResPtr
=
std
::
shared_ptr
<
Resource
>
;
using
ResWPtr
=
std
::
weak_ptr
<
Resource
>
;
class
FaissGpuResourceMgr
{
public:
using
ResBQ
=
zilliz
::
milvus
::
server
::
BlockingQueue
<
ResPtr
>
;
struct
DeviceParams
{
int64_t
temp_mem_size
=
0
;
int64_t
pinned_mem_size
=
0
;
...
...
@@ -55,11 +59,8 @@ class FaissGpuResourceMgr {
// allocate gpu memory before search
// this func will return True if the device is idle and exists an idle resource.
bool
GetRes
(
const
int64_t
&
device_id
,
ResPtr
&
res
,
const
int64_t
&
alloc_size
=
0
);
void
MoveToInuse
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
);
//bool
//GetRes(const int64_t& device_id, ResPtr &res, const int64_t& alloc_size = 0);
void
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
);
...
...
@@ -67,33 +68,34 @@ class FaissGpuResourceMgr {
void
Dump
();
protected:
void
RemoveResource
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
,
std
::
map
<
int64_t
,
std
::
vector
<
ResPtr
>>&
resource_pool
);
protected:
bool
is_init
=
false
;
std
::
mutex
mutex_
;
std
::
map
<
int64_t
,
DeviceParams
>
devices_params_
;
std
::
map
<
int64_t
,
std
::
vector
<
ResPtr
>>
in_use_
;
std
::
map
<
int64_t
,
std
::
vector
<
ResPtr
>>
idle_
;
std
::
map
<
int64_t
,
ResBQ
>
idle_map
;
};
class
ResScope
{
public:
ResScope
(
const
int64_t
device_id
,
ResPtr
&
res
)
:
resource
(
res
),
device_id
(
device_id
)
{
FaissGpuResourceMgr
::
GetInstance
().
MoveToInuse
(
device_id
,
resource
);
ResScope
(
const
int64_t
device_id
,
ResPtr
&
res
)
:
resource
(
res
),
device_id
(
device_id
),
move
(
true
)
{
res
->
mutex
.
lock
();
}
ResScope
(
ResPtr
&
res
)
:
resource
(
res
),
device_id
(
-
1
),
move
(
false
)
{
res
->
mutex
.
lock
();
}
~
ResScope
()
{
//resource->faiss_res->noTempMemory();
FaissGpuResourceMgr
::
GetInstance
().
MoveToIdle
(
device_id
,
resource
);
if
(
move
)
{
FaissGpuResourceMgr
::
GetInstance
().
MoveToIdle
(
device_id
,
resource
);
}
resource
->
mutex
.
unlock
();
}
private:
ResPtr
resource
;
int64_t
device_id
;
bool
move
=
true
;
};
class
GPUIndex
{
...
...
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
浏览文件 @
6dd89bb0
...
...
@@ -130,19 +130,17 @@ void GPUIVF::search_impl(int64_t n,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
)
{
// TODO(linxj): allocate mem
auto
temp_res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
);
if
(
temp_res
)
{
ResScope
rs
(
gpu_id_
,
temp_res
);
if
(
auto
device_index
=
std
::
static_pointer_cast
<
faiss
::
gpu
::
GpuIndexIVF
>
(
index_
))
{
auto
nprobe
=
cfg
.
get_with_default
(
"nprobe"
,
size_t
(
1
));
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
device_index
->
setNumProbes
(
nprobe
);
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
if
(
auto
device_index
=
std
::
static_pointer_cast
<
faiss
::
gpu
::
GpuIndexIVF
>
(
index_
)
)
{
auto
nprobe
=
cfg
.
get_with_default
(
"nprobe"
,
size_t
(
1
)
);
device_index
->
setNumProbes
(
nprobe
);
{
// TODO(linxj): allocate mem
ResScope
rs
(
res_
);
device_index
->
search
(
n
,
(
float
*
)
data
,
k
,
distances
,
labels
);
}
}
else
{
KNOWHERE_THROW_MSG
(
"search can't get gpu resource"
);
}
}
...
...
@@ -283,119 +281,70 @@ void FaissGpuResourceMgr::InitResource() {
is_init
=
true
;
for
(
auto
&
device
:
devices_params_
)
{
auto
&
resource_vec
=
idle_
[
device
.
first
];
auto
&
device_id
=
device
.
first
;
auto
&
device_param
=
device
.
second
;
auto
&
bq
=
idle_map
[
device_id
];
for
(
int64_t
i
=
0
;
i
<
device
.
second
.
resource_num
;
++
i
)
{
auto
r
es
=
std
::
make_shared
<
faiss
::
gpu
::
StandardGpuResources
>
();
for
(
int64_t
i
=
0
;
i
<
device
_param
.
resource_num
;
++
i
)
{
auto
r
aw_resource
=
std
::
make_shared
<
faiss
::
gpu
::
StandardGpuResources
>
();
// TODO(linxj): enable set pinned memory
//res->noTempMemory();
auto
res_wrapper
=
std
::
make_shared
<
Resource
>
(
res
);
AllocateTempMem
(
res_wrapper
,
device
.
first
,
0
);
auto
res_wrapper
=
std
::
make_shared
<
Resource
>
(
raw_resource
);
AllocateTempMem
(
res_wrapper
,
device_id
,
0
);
resource_vec
.
emplace_back
(
res_wrapper
);
bq
.
Put
(
res_wrapper
);
}
}
}
ResPtr
FaissGpuResourceMgr
::
GetRes
(
const
int64_t
&
device_id
,
const
int64_t
&
alloc_size
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
if
(
!
is_init
)
{
InitResource
();
is_init
=
true
;
}
auto
search
=
idle_
.
find
(
device_id
);
if
(
search
!=
idle_
.
end
())
{
auto
res
=
search
->
second
.
back
();
//AllocateTempMem(res, device_id, alloc_size);
search
->
second
.
pop_back
();
return
res
;
InitResource
();
auto
finder
=
idle_map
.
find
(
device_id
);
if
(
finder
!=
idle_map
.
end
())
{
auto
&
bq
=
finder
->
second
;
auto
&&
resource
=
bq
.
Take
();
AllocateTempMem
(
resource
,
device_id
,
alloc_size
);
return
resource
;
}
return
nullptr
;
}
bool
FaissGpuResourceMgr
::
GetRes
(
const
int64_t
&
device_id
,
ResPtr
&
res
,
const
int64_t
&
alloc_size
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
if
(
!
is_init
)
{
InitResource
();
is_init
=
true
;
}
auto
search
=
idle_
.
find
(
device_id
);
if
(
search
!=
idle_
.
end
())
{
auto
&
res_vec
=
search
->
second
;
for
(
auto
it
=
res_vec
.
cbegin
();
it
!=
res_vec
.
cend
();
++
it
)
{
if
((
*
it
)
->
id
==
res
->
id
)
{
//AllocateTempMem(res, device_id, alloc_size);
res_vec
.
erase
(
it
);
return
true
;
}
}
}
// else
return
false
;
}
void
FaissGpuResourceMgr
::
MoveToInuse
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
RemoveResource
(
device_id
,
res
,
idle_
);
in_use_
[
device_id
].
push_back
(
res
);
}
//bool FaissGpuResourceMgr::GetRes(const int64_t &device_id,
// ResPtr &res,
// const int64_t &alloc_size) {
// InitResource();
//
// std::lock_guard<std::mutex> lk(res->mutex);
// AllocateTempMem(res, device_id, alloc_size);
// return true;
//}
void
FaissGpuResourceMgr
::
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
RemoveResource
(
device_id
,
res
,
in_use_
);
auto
it
=
idle_
[
device_id
].
begin
();
idle_
[
device_id
].
insert
(
it
,
res
);
}
void
FaissGpuResourceMgr
::
RemoveResource
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
,
std
::
map
<
int64_t
,
std
::
vector
<
ResPtr
>>
&
resource_pool
)
{
if
(
resource_pool
.
find
(
device_id
)
!=
resource_pool
.
end
())
{
std
::
vector
<
ResPtr
>
&
res_array
=
resource_pool
[
device_id
];
res_array
.
erase
(
std
::
remove_if
(
res_array
.
begin
(),
res_array
.
end
(),
[
&
](
ResPtr
&
ptr
)
{
return
ptr
->
id
==
res
->
id
;
}),
res_array
.
end
());
auto
finder
=
idle_map
.
find
(
device_id
);
if
(
finder
!=
idle_map
.
end
())
{
auto
&
bq
=
finder
->
second
;
bq
.
Put
(
res
);
}
}
void
FaissGpuResourceMgr
::
Free
()
{
for
(
auto
&
item
:
in_use_
)
{
auto
&
res_vec
=
item
.
second
;
res_vec
.
clear
();
}
for
(
auto
&
item
:
idle_
)
{
auto
&
res_vec
=
item
.
second
;
res_vec
.
clear
();
for
(
auto
&
item
:
idle_map
)
{
auto
&
bq
=
item
.
second
;
while
(
!
bq
.
Empty
())
{
bq
.
Take
();
}
}
is_init
=
false
;
}
void
FaissGpuResourceMgr
::
Dump
()
{
std
::
cout
<<
"In used resource"
<<
std
::
endl
;
for
(
auto
&
item
:
in_use_
)
{
std
::
cout
<<
"device_id: "
<<
item
.
first
<<
std
::
endl
;
for
(
auto
&
elem
:
item
.
second
)
{
std
::
cout
<<
"resource_id: "
<<
elem
->
id
<<
std
::
endl
;
}
}
std
::
cout
<<
"Idle resource"
<<
std
::
endl
;
for
(
auto
&
item
:
idle_
)
{
std
::
cout
<<
"device_id: "
<<
item
.
first
<<
std
::
endl
;
for
(
auto
&
elem
:
item
.
second
)
{
std
::
cout
<<
"resource_id: "
<<
elem
->
id
<<
std
::
endl
;
}
for
(
auto
&
item
:
idle_map
)
{
auto
&
bq
=
item
.
second
;
std
::
cout
<<
"device_id: "
<<
item
.
first
<<
", resource count:"
<<
bq
.
Size
();
}
}
...
...
cpp/src/core/test/test_ivf.cpp
浏览文件 @
6dd89bb0
...
...
@@ -386,7 +386,7 @@ class GPURESTEST
int64_t
elems
=
0
;
};
const
int
search_count
=
1
0
;
const
int
search_count
=
1
8
;
const
int
load_count
=
3
;
TEST_F
(
GPURESTEST
,
gpu_ivf_resource_test
)
{
...
...
cpp/src/utils/BlockingQueue.inl
浏览文件 @
6dd89bb0
#pragma once
#include "Log.h"
//
#include "Log.h"
#include "Error.h"
namespace zilliz {
...
...
@@ -17,7 +17,7 @@ BlockingQueue<T>::Put(const T &task) {
std::string error_msg =
"blocking queue is full, capacity: " + std::to_string(capacity_) + " queue_size: " +
std::to_string(queue_.size());
SERVER_LOG_ERROR << error_msg;
//
SERVER_LOG_ERROR << error_msg;
throw ServerException(SERVER_BLOCKING_QUEUE_EMPTY, error_msg);
}
...
...
@@ -33,7 +33,7 @@ BlockingQueue<T>::Take() {
if (queue_.empty()) {
std::string error_msg = "blocking queue empty";
SERVER_LOG_ERROR << error_msg;
//
SERVER_LOG_ERROR << error_msg;
throw ServerException(SERVER_BLOCKING_QUEUE_EMPTY, error_msg);
}
...
...
@@ -57,7 +57,7 @@ BlockingQueue<T>::Front() {
empty_.wait(lock, [this] { return !queue_.empty(); });
if (queue_.empty()) {
std::string error_msg = "blocking queue empty";
SERVER_LOG_ERROR << error_msg;
//
SERVER_LOG_ERROR << error_msg;
throw ServerException(SERVER_BLOCKING_QUEUE_EMPTY, error_msg);
}
T front(queue_.front());
...
...
@@ -72,7 +72,7 @@ BlockingQueue<T>::Back() {
if (queue_.empty()) {
std::string error_msg = "blocking queue empty";
SERVER_LOG_ERROR << error_msg;
//
SERVER_LOG_ERROR << error_msg;
throw ServerException(SERVER_BLOCKING_QUEUE_EMPTY, error_msg);
}
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录