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()来进行异步拷贝的工作。