4. mcVS编程指南

4.1. mcVS索引接口

使用mcVS时，通常遵循的接口调用顺序为：创建索引 -> 添加向量数据 -> 搜索向量。

4.2. 索引的基本使用

本节介绍如何在曦云系列GPU上使用mcVS GPU索引。

4.2.1. 创建索引

创建一个索引类。

代码示例（C++）

使用C++ API时，首先需要创建一个Index类的实例：

raft::device_resources dev_resources;
ivf_flat::index_params index_params;
index_params.n_lists = 1024;
index_params.kmeans_trainset_fraction = 0.1;
index_params.metric = cuvs::distance::DistanceType::L2Expanded;
auto index = ivf_flat::build(dev_resources, index_params, dataset);

示例中，创建了一个GPU索引类，为使用欧式距离（L2）的IVFFlat索引类，nlist为1024。

代码示例（Python）

build_params = ivf_flat.IndexParams(n_lists=1024)
index = ivf_flat.build(build_params, xb)

4.2.2. 添加向量数据

将向量数据添加到索引中。

代码示例（C++）

int nb = 10000;     // 向量数据的数量
auto xb = raft::make_device_matrix<float, int64_t>(dev_resources, nb, d);
// 填充向量数据到xb
index = ivf_flat::extend(dev_resources, xb, std::nullopt, index);

代码示例（Python）

nb = 10000  # 向量数据的数量
# 填充向量数据到xb，这里使用随机数做示例
xb = cp.asarray(np.random.rand(nb, d).astype(np.float32))
index = ivf_flat.extend(index, xb)

添加操作可以反复进行，但是每次添加都有可能触发设备内存的重新分配。因此，为了提高数据添加的效率，建议批量添加数据。

4.2.3. 搜索向量

向量搜索是索引的核心功能，可以通过以下方式进行搜索。

代码示例（C++）

int nq = 10;        // 查询向量的数量
int topk = 3;       // 搜索的topk

auto xq = raft::make_device_matrix<float, int64_t>(dev_resources, n_queries, d);
// 填充查询向量数据到xq

auto neighbors    = raft::make_device_matrix<int64_t>(dev_resources, nq, topk);
auto distances    = raft::make_device_matrix<float>(dev_resources, nq, topk);
ivf_flat::search_params search_params;
search_params.n_probes = 2;
ivf_flat::search(
    dev_resources, search_params, index, xq, neighbors.view(), distances.view());

搜索结束后， distances 中保存了搜索结果的距离，neighbors 中保存了搜索结果的ID。

代码示例（Python）

nq = 10  # 查询向量的数量
topk = 3  # 搜索的topk
# 填充查询向量数据到xq，这里使用随机数做示例
xq = cp.asarray(np.random.rand(nq, d).astype(np.float32))
search_params = ivf_flat.SearchParams()
distances, neighbors = ivf_flat.search(
     search_params,
     index,
     xq,
     topk
)

通常情况下，GPU上的搜索效率会比CPU上的更高，但是在搜索之前需要将查询向量数据拷贝到GPU上，因此在实际使用的过程中，为了提高搜索的效率，建议进行批量搜索。