Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
BaiXuePrincess
milvus
提交
f61aafa9
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,发现更多精彩内容 >>
提交
f61aafa9
编写于
9月 04, 2019
作者:
X
xj.lin
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
MS-476 1. update gpuresource mgr 2. improve gpu index performance
Former-commit-id: 1cdbf8c44da0a5ac88d895cdf29d5aee0a578d59
上级
369d49bc
变更
8
隐藏空白更改
内联
并排
Showing
8 changed file
with
361 addition
and
69 deletion
+361
-69
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
+38
-15
cpp/src/core/include/knowhere/index/vector_index/idmap.h
cpp/src/core/include/knowhere/index/vector_index/idmap.h
+2
-2
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
+94
-42
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
+3
-3
cpp/src/core/src/knowhere/index/vector_index/ivf.cpp
cpp/src/core/src/knowhere/index/vector_index/ivf.cpp
+4
-4
cpp/src/core/test/test_idmap.cpp
cpp/src/core/test/test_idmap.cpp
+2
-0
cpp/src/core/test/test_ivf.cpp
cpp/src/core/test/test_ivf.cpp
+217
-2
cpp/unittest/knowhere/knowhere_test.cpp
cpp/unittest/knowhere/knowhere_test.cpp
+1
-1
未找到文件。
cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
浏览文件 @
f61aafa9
...
...
@@ -8,6 +8,18 @@
namespace
zilliz
{
namespace
knowhere
{
struct
Resource
{
Resource
(
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
&
r
)
:
faiss_res
(
r
)
{
static
int64_t
global_id
=
0
;
id
=
global_id
++
;
}
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
faiss_res
;
int64_t
id
;
};
using
ResPtr
=
std
::
shared_ptr
<
Resource
>
;
using
ResWPtr
=
std
::
weak_ptr
<
Resource
>
;
class
FaissGpuResourceMgr
{
public:
struct
DeviceParams
{
...
...
@@ -17,14 +29,11 @@ class FaissGpuResourceMgr {
};
public:
using
ResPtr
=
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
;
using
ResWPtr
=
std
::
weak_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
;
static
FaissGpuResourceMgr
&
GetInstance
();
void
AllocateTempMem
(
ResPtr
&
res
,
const
int64_t
&
device_id
,
const
int64_t
&
size
);
AllocateTempMem
(
ResPtr
&
res
ource
,
const
int64_t
&
device_id
,
const
int64_t
&
size
);
void
InitDevice
(
int64_t
device_id
,
...
...
@@ -32,12 +41,23 @@ class FaissGpuResourceMgr {
int64_t
temp_mem_size
=
0
,
int64_t
res_num
=
2
);
void
InitResource
();
void
InitResource
();
// allocate gpu memory invoke by build or copy_to_gpu
ResPtr
GetRes
(
const
int64_t
&
device_id
,
const
int64_t
&
alloc_size
=
0
);
ResPtr
GetRes
(
const
int64_t
&
device_id
,
const
int64_t
&
alloc_size
=
0
);
// 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
);
void
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
);
void
MoveToInuse
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
);
void
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
);
protected:
bool
is_init
=
false
;
...
...
@@ -50,23 +70,24 @@ class FaissGpuResourceMgr {
class
ResScope
{
public:
ResScope
(
const
int64_t
device_id
,
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
&
res
)
:
resource
(
res
),
device_id
(
device_id
)
{
ResScope
(
const
int64_t
device_id
,
ResPtr
&
res
)
:
resource
(
res
),
device_id
(
device_id
)
{
FaissGpuResourceMgr
::
GetInstance
().
MoveToInuse
(
device_id
,
resource
);
}
~
ResScope
()
{
resource
->
noTempMemory
();
//resource->faiss_res
->noTempMemory();
FaissGpuResourceMgr
::
GetInstance
().
MoveToIdle
(
device_id
,
resource
);
}
private:
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
resource
;
ResPtr
resource
;
int64_t
device_id
;
};
class
GPUIndex
{
public:
explicit
GPUIndex
(
const
int
&
device_id
)
:
gpu_id_
(
device_id
)
{};
explicit
GPUIndex
(
const
int
&
device_id
)
:
gpu_id_
(
device_id
)
{}
GPUIndex
(
const
int
&
device_id
,
ResPtr
resource
)
:
gpu_id_
(
device_id
),
res_
(
std
::
move
(
resource
)){}
virtual
VectorIndexPtr
CopyGpuToCpu
(
const
Config
&
config
)
=
0
;
virtual
VectorIndexPtr
CopyGpuToGpu
(
const
int64_t
&
device_id
,
const
Config
&
config
)
=
0
;
...
...
@@ -76,13 +97,14 @@ class GPUIndex {
protected:
int64_t
gpu_id_
;
ResPtr
res_
=
nullptr
;
};
class
GPUIVF
:
public
IVF
,
public
GPUIndex
{
public:
explicit
GPUIVF
(
const
int
&
device_id
)
:
IVF
(),
GPUIndex
(
device_id
)
{}
explicit
GPUIVF
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
)
:
IVF
(
std
::
move
(
index
)),
GPUIndex
(
device_id
)
{};
explicit
GPUIVF
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
,
ResPtr
&
resource
)
:
IVF
(
std
::
move
(
index
)),
GPUIndex
(
device_id
,
resource
)
{};
IndexModelPtr
Train
(
const
DatasetPtr
&
dataset
,
const
Config
&
config
)
override
;
void
set_index_model
(
IndexModelPtr
model
)
override
;
//DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override;
...
...
@@ -107,7 +129,8 @@ class GPUIVF : public IVF, public GPUIndex {
class
GPUIVFSQ
:
public
GPUIVF
{
public:
explicit
GPUIVFSQ
(
const
int
&
device_id
)
:
GPUIVF
(
device_id
)
{}
explicit
GPUIVFSQ
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
)
:
GPUIVF
(
std
::
move
(
index
),
device_id
)
{};
explicit
GPUIVFSQ
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
,
ResPtr
&
resource
)
:
GPUIVF
(
std
::
move
(
index
),
device_id
,
resource
)
{};
IndexModelPtr
Train
(
const
DatasetPtr
&
dataset
,
const
Config
&
config
)
override
;
public:
...
...
cpp/src/core/include/knowhere/index/vector_index/idmap.h
浏览文件 @
f61aafa9
...
...
@@ -39,8 +39,8 @@ using IDMAPPtr = std::shared_ptr<IDMAP>;
class
GPUIDMAP
:
public
IDMAP
,
public
GPUIndex
{
public:
explicit
GPUIDMAP
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
)
:
IDMAP
(
std
::
move
(
index
)),
GPUIndex
(
device_id
)
{}
explicit
GPUIDMAP
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
,
ResPtr
&
res
)
:
IDMAP
(
std
::
move
(
index
)),
GPUIndex
(
device_id
,
res
)
{}
VectorIndexPtr
CopyGpuToCpu
(
const
Config
&
config
)
override
;
float
*
GetRawVectors
()
override
;
...
...
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
浏览文件 @
f61aafa9
...
...
@@ -31,26 +31,31 @@ IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) {
GETTENSOR
(
dataset
)
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_device
);
ResScope
rs
(
gpu_device
,
res
);
faiss
::
gpu
::
GpuIndexIVFFlatConfig
idx_config
;
idx_config
.
device
=
gpu_device
;
faiss
::
gpu
::
GpuIndexIVFFlat
device_index
(
res
.
get
(),
dim
,
nlist
,
metric_type
,
idx_config
);
device_index
.
train
(
rows
,
(
float
*
)
p_data
);
std
::
shared_ptr
<
faiss
::
Index
>
host_index
=
nullptr
;
host_index
.
reset
(
faiss
::
gpu
::
index_gpu_to_cpu
(
&
device_index
));
return
std
::
make_shared
<
IVFIndexModel
>
(
host_index
);
auto
temp_resource
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_device
);
if
(
temp_resource
!=
nullptr
)
{
ResScope
rs
(
gpu_device
,
temp_resource
);
faiss
::
gpu
::
GpuIndexIVFFlatConfig
idx_config
;
idx_config
.
device
=
gpu_device
;
faiss
::
gpu
::
GpuIndexIVFFlat
device_index
(
temp_resource
->
faiss_res
.
get
(),
dim
,
nlist
,
metric_type
,
idx_config
);
device_index
.
train
(
rows
,
(
float
*
)
p_data
);
std
::
shared_ptr
<
faiss
::
Index
>
host_index
=
nullptr
;
host_index
.
reset
(
faiss
::
gpu
::
index_gpu_to_cpu
(
&
device_index
));
return
std
::
make_shared
<
IVFIndexModel
>
(
host_index
);
}
else
{
KNOWHERE_THROW_MSG
(
"Build IVF can't get gpu resource"
);
}
}
void
GPUIVF
::
set_index_model
(
IndexModelPtr
model
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
auto
host_index
=
std
::
static_pointer_cast
<
IVFIndexModel
>
(
model
);
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
))
{
ResScope
rs
(
gpu_id_
,
res
);
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
gpu_id_
,
host_index
->
index_
.
get
());
if
(
auto
gpures
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
))
{
ResScope
rs
(
gpu_id_
,
gpures
);
res_
=
gpures
;
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res_
->
faiss_res
.
get
(),
gpu_id_
,
host_index
->
index_
.
get
());
index_
.
reset
(
device_index
);
}
else
{
KNOWHERE_THROW_MSG
(
"load index model error, can't get gpu_resource"
);
...
...
@@ -94,9 +99,10 @@ void GPUIVF::LoadImpl(const BinarySet &index_binary) {
faiss
::
Index
*
index
=
faiss
::
read_index
(
&
reader
);
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
))
{
ResScope
rs
(
gpu_id_
,
res
);
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
gpu_id_
,
index
);
if
(
auto
temp_res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
))
{
ResScope
rs
(
gpu_id_
,
temp_res
);
res_
=
temp_res
;
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res_
->
faiss_res
.
get
(),
gpu_id_
,
index
);
index_
.
reset
(
device_index
);
}
else
{
KNOWHERE_THROW_MSG
(
"Load error, can't get gpu resource"
);
...
...
@@ -123,14 +129,21 @@ void GPUIVF::search_impl(int64_t n,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
)
{
if
(
auto
device_index
=
std
::
static_pointer_cast
<
faiss
::
gpu
::
GpuIndexIVF
>
(
index_
))
{
// todo: allocate search memory
auto
nprobe
=
cfg
.
get_with_default
(
"nprobe"
,
size_t
(
1
));
// TODO(linxj): allocate mem
if
(
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
,
res_
))
{
ResScope
rs
(
gpu_id_
,
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
);
device_index
->
search
(
n
,
(
float
*
)
data
,
k
,
distances
,
labels
);
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
device_index
->
setNumProbes
(
nprobe
);
device_index
->
search
(
n
,
(
float
*
)
data
,
k
,
distances
,
labels
);
}
}
else
{
KNOWHERE_THROW_MSG
(
"search can't get gpu resource"
);
}
}
VectorIndexPtr
GPUIVF
::
CopyGpuToCpu
(
const
Config
&
config
)
{
...
...
@@ -165,6 +178,7 @@ IndexModelPtr GPUIVFPQ::Train(const DatasetPtr &dataset, const Config &config) {
GETTENSOR
(
dataset
)
// TODO(linxj): set device here.
// TODO(linxj): set gpu resource here.
faiss
::
gpu
::
StandardGpuResources
res
;
faiss
::
gpu
::
GpuIndexIVFPQ
device_index
(
&
res
,
dim
,
nlist
,
M
,
nbits
,
metric_type
);
device_index
.
train
(
rows
,
(
float
*
)
p_data
);
...
...
@@ -202,17 +216,23 @@ IndexModelPtr GPUIVFSQ::Train(const DatasetPtr &dataset, const Config &config) {
index_type
<<
"IVF"
<<
nlist
<<
","
<<
"SQ"
<<
nbits
;
auto
build_index
=
faiss
::
index_factory
(
dim
,
index_type
.
str
().
c_str
(),
metric_type
);
faiss
::
gpu
::
StandardGpuResources
res
;
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
&
res
,
gpu_num
,
build_index
);
device_index
->
train
(
rows
,
(
float
*
)
p_data
);
auto
temp_resource
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_num
);
if
(
temp_resource
!=
nullptr
)
{
ResScope
rs
(
gpu_num
,
temp_resource
);
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
temp_resource
->
faiss_res
.
get
(),
gpu_num
,
build_index
);
device_index
->
train
(
rows
,
(
float
*
)
p_data
);
std
::
shared_ptr
<
faiss
::
Index
>
host_index
=
nullptr
;
host_index
.
reset
(
faiss
::
gpu
::
index_gpu_to_cpu
(
device_index
));
std
::
shared_ptr
<
faiss
::
Index
>
host_index
=
nullptr
;
host_index
.
reset
(
faiss
::
gpu
::
index_gpu_to_cpu
(
device_index
));
delete
device_index
;
delete
build_index
;
delete
device_index
;
delete
build_index
;
return
std
::
make_shared
<
IVFIndexModel
>
(
host_index
);
return
std
::
make_shared
<
IVFIndexModel
>
(
host_index
);
}
else
{
KNOWHERE_THROW_MSG
(
"Build IVFSQ can't get gpu resource"
);
}
}
VectorIndexPtr
GPUIVFSQ
::
CopyGpuToCpu
(
const
Config
&
config
)
{
...
...
@@ -231,16 +251,16 @@ FaissGpuResourceMgr &FaissGpuResourceMgr::GetInstance() {
return
instance
;
}
void
FaissGpuResourceMgr
::
AllocateTempMem
(
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
&
res
,
void
FaissGpuResourceMgr
::
AllocateTempMem
(
ResPtr
&
resource
,
const
int64_t
&
device_id
,
const
int64_t
&
size
)
{
if
(
size
)
{
res
->
setTempMemory
(
size
);
res
ource
->
faiss_res
->
setTempMemory
(
size
);
}
else
{
auto
search
=
devices_params_
.
find
(
device_id
);
if
(
search
!=
devices_params_
.
end
())
{
res
->
setTempMemory
(
search
->
second
.
temp_mem_size
);
res
ource
->
faiss_res
->
setTempMemory
(
search
->
second
.
temp_mem_size
);
}
// else do nothing. allocate when use.
}
...
...
@@ -264,14 +284,19 @@ void FaissGpuResourceMgr::InitResource() {
for
(
int
i
=
0
;
i
<
device
.
second
.
resource_num
;
++
i
)
{
auto
res
=
std
::
make_shared
<
faiss
::
gpu
::
StandardGpuResources
>
();
res
->
noTempMemory
();
resource_vec
.
push_back
(
res
);
// TODO(linxj): enable set pinned memory
//res->noTempMemory();
auto
res_wrapper
=
std
::
make_shared
<
Resource
>
(
res
);
AllocateTempMem
(
res_wrapper
,
device
.
first
,
0
);
resource_vec
.
emplace_back
(
res_wrapper
);
}
}
}
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
FaissGpuResourceMgr
::
GetRes
(
const
int64_t
&
device_id
,
const
int64_t
&
alloc_size
)
{
ResPtr
FaissGpuResourceMgr
::
GetRes
(
const
int64_t
&
device_id
,
const
int64_t
&
alloc_size
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
if
(
!
is_init
)
{
...
...
@@ -282,21 +307,48 @@ std::shared_ptr<faiss::gpu::StandardGpuResources> FaissGpuResourceMgr::GetRes(co
auto
search
=
idle_
.
find
(
device_id
);
if
(
search
!=
idle_
.
end
())
{
auto
res
=
search
->
second
.
back
();
AllocateTempMem
(
res
,
device_id
,
alloc_size
);
//
AllocateTempMem(res, device_id, alloc_size);
search
->
second
.
pop_back
();
return
res
;
}
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
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
&
res
)
{
void
FaissGpuResourceMgr
::
MoveToInuse
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
in_use_
[
device_id
].
push_back
(
res
);
}
void
FaissGpuResourceMgr
::
MoveToIdle
(
const
int64_t
&
device_id
,
const
std
::
shared_ptr
<
faiss
::
gpu
::
StandardGpuResources
>
&
res
)
{
void
FaissGpuResourceMgr
::
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
idle_
[
device_id
].
push_back
(
res
);
auto
it
=
idle_
[
device_id
].
begin
();
idle_
[
device_id
].
insert
(
it
,
res
);
}
void
GPUIndex
::
SetGpuDevice
(
const
int
&
gpu_id
)
{
...
...
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
浏览文件 @
f61aafa9
...
...
@@ -135,11 +135,11 @@ VectorIndexPtr IDMAP::Clone() {
VectorIndexPtr
IDMAP
::
CopyCpuToGpu
(
const
int64_t
&
device_id
,
const
Config
&
config
)
{
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
device_id
)){
ResScope
rs
(
device_id
,
res
);
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
device_id
,
index_
.
get
());
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
device_id
,
index_
.
get
());
std
::
shared_ptr
<
faiss
::
Index
>
device_index
;
device_index
.
reset
(
gpu_index
);
return
std
::
make_shared
<
GPUIDMAP
>
(
device_index
,
device_id
);
return
std
::
make_shared
<
GPUIDMAP
>
(
device_index
,
device_id
,
res
);
}
else
{
KNOWHERE_THROW_MSG
(
"CopyCpuToGpu Error, can't get gpu_resource"
);
}
...
...
@@ -204,7 +204,7 @@ void GPUIDMAP::LoadImpl(const BinarySet &index_binary) {
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
)
){
ResScope
rs
(
gpu_id_
,
res
);
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
gpu_id_
,
index
);
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
gpu_id_
,
index
);
index_
.
reset
(
device_index
);
}
else
{
KNOWHERE_THROW_MSG
(
"Load error, can't get gpu resource"
);
...
...
cpp/src/core/src/knowhere/index/vector_index/ivf.cpp
浏览文件 @
f61aafa9
...
...
@@ -197,11 +197,11 @@ void IVF::search_impl(int64_t n,
VectorIndexPtr
IVF
::
CopyCpuToGpu
(
const
int64_t
&
device_id
,
const
Config
&
config
)
{
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
device_id
)){
ResScope
rs
(
device_id
,
res
);
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
device_id
,
index_
.
get
());
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
device_id
,
index_
.
get
());
std
::
shared_ptr
<
faiss
::
Index
>
device_index
;
device_index
.
reset
(
gpu_index
);
return
std
::
make_shared
<
GPUIVF
>
(
device_index
,
device_id
);
return
std
::
make_shared
<
GPUIVF
>
(
device_index
,
device_id
,
res
);
}
else
{
KNOWHERE_THROW_MSG
(
"CopyCpuToGpu Error, can't get gpu_resource"
);
}
...
...
@@ -275,11 +275,11 @@ VectorIndexPtr IVFSQ::CopyCpuToGpu(const int64_t &device_id, const Config &confi
faiss
::
gpu
::
GpuClonerOptions
option
;
option
.
allInGpu
=
true
;
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
.
get
(),
device_id
,
index_
.
get
(),
&
option
);
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
device_id
,
index_
.
get
(),
&
option
);
std
::
shared_ptr
<
faiss
::
Index
>
device_index
;
device_index
.
reset
(
gpu_index
);
return
std
::
make_shared
<
GPUIVFSQ
>
(
device_index
,
device_id
);
return
std
::
make_shared
<
GPUIVFSQ
>
(
device_index
,
device_id
,
res
);
}
else
{
KNOWHERE_THROW_MSG
(
"CopyCpuToGpu Error, can't get gpu_resource"
);
}
...
...
cpp/src/core/test/test_idmap.cpp
浏览文件 @
f61aafa9
...
...
@@ -18,9 +18,11 @@
using
namespace
zilliz
::
knowhere
;
static
int
device_id
=
0
;
class
IDMAPTest
:
public
DataGen
,
public
::
testing
::
Test
{
protected:
void
SetUp
()
override
{
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
device_id
,
1024
*
1024
*
200
,
1024
*
1024
*
300
,
2
);
Init_with_default
();
index_
=
std
::
make_shared
<
IDMAP
>
();
}
...
...
cpp/src/core/test/test_ivf.cpp
浏览文件 @
f61aafa9
...
...
@@ -8,6 +8,12 @@
#include <iostream>
#include <sstream>
#include <thread>
#include <faiss/AutoTune.h>
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexIVFFlat.h>
#include <faiss/gpu/GpuClonerOptions.h>
#include "knowhere/index/vector_index/gpu_ivf.h"
#include "knowhere/index/vector_index/ivf.h"
...
...
@@ -25,7 +31,7 @@ using ::testing::TestWithParam;
using
::
testing
::
Values
;
using
::
testing
::
Combine
;
static
int
device_id
=
1
;
static
int
device_id
=
0
;
IVFIndexPtr
IndexFactory
(
const
std
::
string
&
type
)
{
if
(
type
==
"IVF"
)
{
return
std
::
make_shared
<
IVF
>
();
...
...
@@ -50,7 +56,7 @@ class IVFTest
//Init_with_default();
Generate
(
128
,
1000000
/
5
,
10
);
index_
=
IndexFactory
(
index_type
);
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
device_id
);
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
device_id
,
1024
*
1024
*
200
,
1024
*
1024
*
300
,
2
);
}
protected:
...
...
@@ -343,4 +349,213 @@ TEST_P(IVFTest, seal_test) {
ASSERT_GE
(
without_seal
,
with_seal
);
}
class
GPURESTEST
:
public
DataGen
,
public
::
testing
::
Test
{
protected:
void
SetUp
()
override
{
//std::tie(index_type, preprocess_cfg, train_cfg, add_cfg, search_cfg) = GetParam();
//Init_with_default();
Generate
(
128
,
1000000
,
1000
);
k
=
100
;
//index_ = IndexFactory(index_type);
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
device_id
,
1024
*
1024
*
200
,
1024
*
1024
*
300
,
2
);
elems
=
nq
*
k
;
ids
=
(
int64_t
*
)
malloc
(
sizeof
(
int64_t
)
*
elems
);
dis
=
(
float
*
)
malloc
(
sizeof
(
float
)
*
elems
);
}
void
TearDown
()
override
{
delete
ids
;
delete
dis
;
}
protected:
std
::
string
index_type
;
Config
preprocess_cfg
;
Config
train_cfg
;
Config
add_cfg
;
Config
search_cfg
;
IVFIndexPtr
index_
=
nullptr
;
int64_t
*
ids
=
nullptr
;
float
*
dis
=
nullptr
;
int64_t
elems
=
0
;
};
const
int
search_count
=
100
;
const
int
load_count
=
30
;
TEST_F
(
GPURESTEST
,
gpu_ivf_resource_test
)
{
assert
(
!
xb
.
empty
());
{
index_type
=
"GPUIVF"
;
index_
=
IndexFactory
(
index_type
);
auto
preprocessor
=
index_
->
BuildPreprocessor
(
base_dataset
,
preprocess_cfg
);
index_
->
set_preprocessor
(
preprocessor
);
train_cfg
=
Config
::
object
{{
"nlist"
,
1638
},
{
"gpu_id"
,
device_id
},
{
"metric_type"
,
"L2"
}};
auto
model
=
index_
->
Train
(
base_dataset
,
train_cfg
);
index_
->
set_index_model
(
model
);
index_
->
Add
(
base_dataset
,
add_cfg
);
EXPECT_EQ
(
index_
->
Count
(),
nb
);
EXPECT_EQ
(
index_
->
Dimension
(),
dim
);
search_cfg
=
Config
::
object
{{
"k"
,
k
}};
TimeRecorder
tc
(
"knowere GPUIVF"
);
for
(
int
i
=
0
;
i
<
search_count
;
++
i
)
{
index_
->
Search
(
query_dataset
,
search_cfg
);
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSection
(
"search all"
);
}
{
// IVF-Search
faiss
::
gpu
::
StandardGpuResources
res
;
faiss
::
gpu
::
GpuIndexIVFFlatConfig
idx_config
;
idx_config
.
device
=
device_id
;
faiss
::
gpu
::
GpuIndexIVFFlat
device_index
(
&
res
,
dim
,
1638
,
faiss
::
METRIC_L2
,
idx_config
);
device_index
.
train
(
nb
,
xb
.
data
());
device_index
.
add
(
nb
,
xb
.
data
());
TimeRecorder
tc
(
"ori IVF"
);
for
(
int
i
=
0
;
i
<
search_count
;
++
i
)
{
device_index
.
search
(
nq
,
xq
.
data
(),
k
,
dis
,
ids
);
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSection
(
"search all"
);
}
}
TEST_F
(
GPURESTEST
,
gpuivfsq
)
{
{
// knowhere gpu ivfsq
index_type
=
"GPUIVFSQ"
;
index_
=
IndexFactory
(
index_type
);
auto
preprocessor
=
index_
->
BuildPreprocessor
(
base_dataset
,
preprocess_cfg
);
index_
->
set_preprocessor
(
preprocessor
);
train_cfg
=
Config
::
object
{{
"gpu_id"
,
device_id
},
{
"nlist"
,
1638
},
{
"nbits"
,
8
},
{
"metric_type"
,
"L2"
}};
auto
model
=
index_
->
Train
(
base_dataset
,
train_cfg
);
index_
->
set_index_model
(
model
);
index_
->
Add
(
base_dataset
,
add_cfg
);
search_cfg
=
Config
::
object
{{
"k"
,
k
}};
auto
result
=
index_
->
Search
(
query_dataset
,
search_cfg
);
AssertAnns
(
result
,
nq
,
k
);
auto
cpu_idx
=
CopyGpuToCpu
(
index_
,
Config
());
cpu_idx
->
Seal
();
TimeRecorder
tc
(
"knowhere GPUSQ8"
);
auto
search_idx
=
CopyCpuToGpu
(
cpu_idx
,
device_id
,
Config
());
tc
.
RecordSection
(
"Copy to gpu"
);
for
(
int
i
=
0
;
i
<
search_count
;
++
i
)
{
search_idx
->
Search
(
query_dataset
,
search_cfg
);
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSection
(
"search all"
);
}
{
// Ori gpuivfsq Test
const
char
*
index_description
=
"IVF1638,SQ8"
;
faiss
::
Index
*
ori_index
=
faiss
::
index_factory
(
dim
,
index_description
,
faiss
::
METRIC_L2
);
faiss
::
gpu
::
StandardGpuResources
res
;
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
&
res
,
device_id
,
ori_index
);
device_index
->
train
(
nb
,
xb
.
data
());
device_index
->
add
(
nb
,
xb
.
data
());
auto
cpu_index
=
faiss
::
gpu
::
index_gpu_to_cpu
(
device_index
);
auto
idx
=
dynamic_cast
<
faiss
::
IndexIVF
*>
(
cpu_index
);
if
(
idx
!=
nullptr
)
{
idx
->
to_readonly
();
}
delete
device_index
;
delete
ori_index
;
faiss
::
gpu
::
GpuClonerOptions
option
;
option
.
allInGpu
=
true
;
TimeRecorder
tc
(
"ori GPUSQ8"
);
faiss
::
Index
*
search_idx
=
faiss
::
gpu
::
index_cpu_to_gpu
(
&
res
,
device_id
,
cpu_index
,
&
option
);
tc
.
RecordSection
(
"Copy to gpu"
);
for
(
int
i
=
0
;
i
<
search_count
;
++
i
)
{
search_idx
->
search
(
nq
,
xq
.
data
(),
k
,
dis
,
ids
);
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSection
(
"search all"
);
delete
cpu_index
;
delete
search_idx
;
}
}
TEST_F
(
GPURESTEST
,
copyandsearch
)
{
printf
(
"==================
\n
"
);
// search and copy at the same time
index_type
=
"GPUIVFSQ"
;
//index_type = "GPUIVF";
index_
=
IndexFactory
(
index_type
);
auto
preprocessor
=
index_
->
BuildPreprocessor
(
base_dataset
,
preprocess_cfg
);
index_
->
set_preprocessor
(
preprocessor
);
train_cfg
=
Config
::
object
{{
"gpu_id"
,
device_id
},
{
"nlist"
,
1638
},
{
"nbits"
,
8
},
{
"metric_type"
,
"L2"
}};
auto
model
=
index_
->
Train
(
base_dataset
,
train_cfg
);
index_
->
set_index_model
(
model
);
index_
->
Add
(
base_dataset
,
add_cfg
);
search_cfg
=
Config
::
object
{{
"k"
,
k
}};
auto
result
=
index_
->
Search
(
query_dataset
,
search_cfg
);
AssertAnns
(
result
,
nq
,
k
);
auto
cpu_idx
=
CopyGpuToCpu
(
index_
,
Config
());
cpu_idx
->
Seal
();
auto
search_idx
=
CopyCpuToGpu
(
cpu_idx
,
device_id
,
Config
());
auto
search_func
=
[
&
]
{
//TimeRecorder tc("search&load");
for
(
int
i
=
0
;
i
<
search_count
;
++
i
)
{
search_idx
->
Search
(
query_dataset
,
search_cfg
);
//if (i > search_count - 6 || i == 0)
// tc.RecordSection("search once");
}
//tc.ElapseFromBegin("search finish");
};
auto
load_func
=
[
&
]
{
//TimeRecorder tc("search&load");
for
(
int
i
=
0
;
i
<
load_count
;
++
i
)
{
CopyCpuToGpu
(
cpu_idx
,
device_id
,
Config
());
//if (i > load_count -5 || i < 5)
//tc.RecordSection("Copy to gpu");
}
//tc.ElapseFromBegin("load finish");
};
TimeRecorder
tc
(
"basic"
);
CopyCpuToGpu
(
cpu_idx
,
device_id
,
Config
());
tc
.
RecordSection
(
"Copy to gpu once"
);
search_idx
->
Search
(
query_dataset
,
search_cfg
);
tc
.
RecordSection
(
"search once"
);
search_func
();
tc
.
RecordSection
(
"only search total"
);
load_func
();
tc
.
RecordSection
(
"only copy total"
);
std
::
thread
search_thread
(
search_func
);
std
::
thread
load_thread
(
load_func
);
search_thread
.
join
();
load_thread
.
join
();
tc
.
RecordSection
(
"Copy&search total"
);
}
// TODO(linxj): Add exception test
cpp/unittest/knowhere/knowhere_test.cpp
浏览文件 @
f61aafa9
...
...
@@ -29,7 +29,7 @@ class KnowhereWrapperTest
:
public
TestWithParam
<::
std
::
tuple
<
IndexType
,
std
::
string
,
int
,
int
,
int
,
int
,
Config
,
Config
>>
{
protected:
void
SetUp
()
override
{
zilliz
::
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
DEVICE_ID
);
zilliz
::
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
DEVICE_ID
,
1024
*
1024
*
200
,
1024
*
1024
*
300
,
2
);
std
::
string
generator_type
;
std
::
tie
(
index_type
,
generator_type
,
dim
,
nb
,
nq
,
k
,
train_cfg
,
search_cfg
)
=
GetParam
();
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录