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: 原向量个数，这里为1k: 要查询的近邻个数，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()来进行异步拷贝的工作。