Faiss(9)：将Index从CPU复制到GPU过程分析

1. 说明

在前面的文章里提到使用了一个python语言编写的APP，用于测试faiss 的search时间，在该程序中train和add Index的过程都是在CPU中进行的，而search的过程则在GPU中进行，这就涉及到将Index从CPU的内存中复制到GPU的内存的过程。

个人对这一过程如何实现颇感兴趣，所以特地对此进行分析。

2. 过程分析

2.1 应用程序

假设此时已经创建了一个可用的Index，已经train，add过了等等。
在程序中调用：

gpu_index = faiss.index_cpu_to_all_gpus(index)

其中index是已经创建的IndexIVFPQ的实例，返回值gpu_index是该Index在GPU中的实例。

2.2 faiss core (python)

faiss源代码编译成python的接口全部定义在python/faiss.py中。

index_cpu_to_all_gpus函数

def index_cpu_to_all_gpus(index, co=None, ngpu=-1):
    if ngpu == -1:
        ngpu = get_num_gpus()
    res = [StandardGpuResources() for i in range(ngpu)]
    index2 = index_cpu_to_gpu_multiple_py(res, index, co)
    return index2

其中：

• index: 上一层传递进来的索引实例
• co : None

这个方法简单地获取了一下当前系统中GPU显卡的个数，并将其作为资源传递给下一级方法。

index_cpu_to_gpu_multiple_py函数
这个方法为GPU索引建立C++向量和资源。并将资源分配给第一个gpu。

def index_cpu_to_gpu_multiple_py(resources, index, co=None):
    """builds the C++ vectors for the GPU indices and the
    resources. Handles the common case where the resources are assigned to
    the first len(resources) GPUs"""
    vres = GpuResourcesVector()
    vdev = IntVector()
    for i, res in enumerate(resources):
        vdev.push_back(i)
        vres.push_back(res)
    index = index_cpu_to_gpu_multiple(vres, vdev, index, co)
    index.referenced_objects = resources
    return index

其中：

• res：一组gpu资源的列表，包含了每一个gpu的资源
• index: 上一层传递进来的索引实例
• co : None

这个方法调用的关键函数index_cpu_to_gpu_multiple()是C++提供的一个API接口。

2.3 faiss core(C++)

index_cpu_to_gpu_multiple函数
这个函数定义在gpu/GpuCloner.cpp文件中，该文件主要定义了CPU与GPU之间复制索引的基本方法。其中index_cpu_gpu_multiple函数定义如下：

// 此函数将index从cpu复制到多个gpu中
// resource: gpu资源容器
// devices: gpu编号，与容器内资源一一对应
// index: 要复制的索引
// options: 复制时的参数，传进来为NULL
faiss::Index * index_cpu_to_gpu_multiple(
       std::vector<GpuResources*> & resources,
       std::vector<int> &devices,
       const faiss::Index *index,
       const GpuMultipleClonerOptions *options)
{
    GpuMultipleClonerOptions defaults;
    ToGpuClonerMultiple cl(resources, devices, options ? *options : defaults);
    return cl.clone_Index(index);
}

函数内定义了一个ToGpuClonerMultiple 类型的结构体实例cl，该结构体包含一个sub_cloners的容器，在构造函数中依次将所有的resource根据options复制到sub_cloners的容器内，然后调用clone_Index()函数来复制索引。

所谓资源，是一组GpuResource结构体的实例，该结构体定义了一组关于GPU处理streams的函数，定义在gpu/GpuResource.h中。

clone_Index()函数
这是复制Index整体到所有的GPU中的函数定义，在它内部还是会调用到针对每一个GPU的克隆函数。

Index *ToGpuClonerMultiple::clone_Index(const Index *index)
{
    // 如果只有一个sub_cloners，那么只需在ToGpuCloner中调用一次即可。
    long n = sub_cloners.size();
    if (n == 1)
        return sub_cloners[0].clone_Index(index);

    if(dynamic_cast<const IndexFlat *>(index) ||
       dynamic_cast<const faiss::IndexIVFFlat *>(index) ||
       dynamic_cast<const faiss::IndexIVFScalarQuantizer *>(index) ||
       dynamic_cast<const faiss::IndexIVFPQ *>(index)) {
        if(!shard) {
            IndexReplicas * res = new IndexReplicas();
            for(auto & sub_cloner: sub_cloners) {
                res->addIndex(sub_cloner.clone_Index(index));
            }
            res->own_fields = true;
            return res;
        } else {
            return clone_Index_to_shards (index);
        }
    } 
    //这里只保留了IVFPQ类型的索引的操作
    ...
}

shard是GpuMultipleClonerOptions 结构体中的一个标志符号，为真时表示跨gpu分割索引，查询会在每个gpu的索引中进行，完成后再将结果合并起来。这里的每一个索引都由一个单独的CPU线程管理。为假时表示跨gpu复制索引。

由于这里使用default的options，而默认shard标志位为false，所以这里会运行前面的if语句中的内容。

注：这里以多个GPU的情况为例。如果只有一个GPU，那么直接调用sub_cloners[0].clone_Index(index);后续操作更简单。

IndexReplicas()是一个模板类，该类针对不同的Index类型创建不同的类，在构造函数中会创建ThreadedIndex的模板类实例，该实例中包含一些针对Index的函数如：addIndex, removeIndex, onAfterAddIndex等等。定义在impl/ThreadedIndex-inl.h文件中。

addIndex()函数的参数是sub-cloner的GPU克隆的Index, 即Index的分片。在addIndex()中会针对每一个这样的分片创建一个线程来管理该Index分片。

clone_Index()函数
这是sub-cloner的函数定义，即会对Index分片后克隆到每一个GPU中。

Index *ToGpuCloner::clone_Index(const Index *index)
{
    //这里只保留IVFPQ类型的索引部分
    ...
    else if(auto ipq = dynamic_cast<const faiss::IndexIVFPQ *>(index)) {
        if(verbose)
            printf("  IndexIVFPQ size %ld -> GpuIndexIVFPQ "
                   "indicesOptions=%d "
                   "usePrecomputed=%d useFloat16=%d reserveVecs=%ld\n",
                   ipq->ntotal, indicesOptions, usePrecomputed,
                   useFloat16, reserveVecs);
        GpuIndexIVFPQConfig config;
        config.device = device;
        config.indicesOptions = indicesOptions;
        config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
        config.flatConfig.storeTransposed = storeTransposed;
        config.useFloat16LookupTables = useFloat16;
        config.usePrecomputedTables = usePrecomputed;

        GpuIndexIVFPQ *res = new GpuIndexIVFPQ(resources, ipq, config);

        if(reserveVecs > 0 && ipq->ntotal == 0) {
            res->reserveMemory(reserveVecs);
        }

        return res;
    } 
    ...
}

在这部分内容中会分别创建GpuIndexIVFPQConfig和GpuIndexIVFPQ类型的实例，这两个类是在cuda类型的源文件中定义的，也就是说这是在GPU中定义和分配空间的变量，然后根据原有Index的config和实例组合成Gpu下同类型的实例。返回值res是Index在GPU的内存空间的首地址，后续res调用的操作都是在GPU中完成。

至此，便完成了Index从CPU复制GPU的全过程。

注：这个函数的返回值会作为参数传递给ThreadedIndex结构体的addIndex函数，也就是说无论Index在GPU中分成了多少个分片，在CPU中都会有一个线程池来一一管理这些GPU的索引。虽然Index内的诸如train, add, search等操作都在GPU中进行，但是CPU中仍然保留了对Index的基本控制能力，如添加Index，移除Index等的操作。

3. 其他关键分支

clone_Index_to_shards()函数