Faiss(11)：关于GPU-CPU的数据Copy back

1. 前言

在之前的文档中记录了Faiss框架search时各个阶段的逻辑顺序和时间消耗，其中发现在第2.3节GPUIndex的搜索中，Copy back占的时间比值不小（达到了45.61%），相信如果要对整体方案进行优化，那么这一部分将是一个重要的突破口。

所以这篇文档主要对数据的copy back进行分析。

2. Copy back说明

GPU search完成后会将输出结果distances和labels的数组保存到内存中，但仍然是在GPU内部的内存，应用程序无法知道实际结果，那么就需要将这部分数据从GPU拷贝到CPU的内存中，然后返回给应用程序。

整个拷贝分成两部分数据：

• Distances是搜索后的topK个近邻与目标向量的距离数组，由小到大排列，float型；
• Labels是上述topK个近邻的向量标签，与Distances一一对应，Index::idx_t型，在本机中是uint64。
  
  

3. 代码追踪

Distances和Labels是调用的同一个函数，即入口相同，只是因为数据类型和地址空间不一样进行了两次拷贝，所以这里只分析Distances的拷贝，Labels完全一样。

3.1 调用入口

调用入口在gpu/GpuIndex.cu文件的GpuIndex::search函数末尾，此时已经完成了search，准备把结果拷贝回CPU。

fromDevice<float, 2>(outDistances, distances, stream);

• outDistances: GPU内为存储结果距离向量而分配的内存，DeviceTensor指针；
• distances: CPU内为存储距离结果而分配的内存，float*类型；
• stream: 对应GPU设备的所有用于此索引的计算流，int型;

参数outDistances

outDistances是GPU内分配的用于存放距离向量的空间，是一个模板类的实例，定义在gpu/GpuIndex.cu中。

定义如下：

auto outDistances =
    toDevice<float, 2>(resources_, device_, distances, stream,
                       {(int) n, (int) k});

• resources_是GpuIndex类的保护型变量，GpuResources类的实例，用于管理streams, cuBLAS句柄和设备内存；
• device_是一个int型变量，用于指定当前所处的GPU；
• distances是float型的指针，指向CPU中为结果分配的首地址；
• stream：对应GPU设备的所有用于此索引的计算流，int型；
• n: 原向量个数，这里为1
• k: 要查询的近邻个数，100

toDevice定义

toDevice是一个函数模板，定义在gpu/utils/CopyUtils.cuh文件中。

template <typename T, int Dim>
DeviceTensor<T, Dim, true> toDevice(GpuResources* resources,
                                    int dstDevice,
                                    T* src,
                                    cudaStream_t stream,
                                    std::initializer_list<int> sizes) {
  int dev = getDeviceForAddress(src);

  if (dev == dstDevice) {
    // On device we expect
    return DeviceTensor<T, Dim, true>(src, sizes);
  } else {
    // On different device or on host
    DeviceScope scope(dstDevice);

    Tensor<T, Dim, true> oldT(src, sizes);

    if (resources) {
      DeviceTensor<T, Dim, true> newT(resources->getMemoryManager(dstDevice),
                                      sizes,
                                      stream);

      newT.copyFrom(oldT, stream);
      return newT;
    } else {
      DeviceTensor<T, Dim, true> newT(sizes);

      newT.copyFrom(oldT, stream);
      return newT;
    }
  }
}

从代码中可以看到，该函数首先判断src给定的地址是否在GPU设备上，是则直接返回一个DeviceTensor类型的指针，否则需要将其复制到GPU上。
实际使用的程序提供了"resource"，那么还需要临时分配内存。


DeviceTensor在本程序中的定义如下(DeviceTensor.cuh)：

__host__ DeviceTensor(DeviceMemory& m,
                        std::initializer_list<IndexT> sizes,
                        cudaStream_t stream,
                        MemorySpace space = MemorySpace::Device);

__host__是CUDA编程定义的声明符，表示该函数在主机上执行或者仅可通过主机调用。

3.2 fromDevice定义

/// Copies a device array's allocation to an address, if necessary
template <typename T>
inline void fromDevice(T* src, T* dst, size_t num, cudaStream_t stream) {
  // 如果目标地址和源地址相同，则无需复制
  if (src == dst) {
    return;
  }

  int dev = getDeviceForAddress(dst);

  // dev==-1，表示目标地址在CPU中
  if (dev == -1) {
    CUDA_VERIFY(cudaMemcpyAsync(dst,
                                src,
                                num * sizeof(T),
                                cudaMemcpyDeviceToHost,
                                stream));
  } else {
    CUDA_VERIFY(cudaMemcpyAsync(dst,
                                src,
                                num * sizeof(T),
                                cudaMemcpyDeviceToDevice,
                                stream));
  }
}

/// Copies a device array's allocation to an address, if necessary
template <typename T, int Dim>
void fromDevice(Tensor<T, Dim, true>& src, T* dst, cudaStream_t stream) {
  FAISS_ASSERT(src.isContiguous());
  fromDevice(src.data(), dst, src.numElements(), stream);
}

• src是GPU的内存地址，dst是CPU的内存地址，函数的作用是将src的data拷贝到dst中。

• CUDA_VERIFY是一个宏函数，用于判断CUDA执行结果是否为cudaSuccess。

根据目标地址所处的位置，判断数据是从Device拷贝到Host中还是从Device拷贝到另一个Device中。然后调用cudaMemcpyAsync()来进行异步拷贝的工作。