kNN-CUDA 源码分析

kNN-CUDA 源码分析

3140105324
文家豪

---

kNN-CUDA 是一个用 CUDA 实现的计算 k 近邻搜索的项目。这个项目用 GPGPU 实现了有两种暴力的算法——它们通过 NVIDA 的 API 在显卡上并行的求解 kNN 搜索。

1. 问题说明

kNN 搜索是指通过计算 query 的向量和 reference 集合内的向量的相似度，并找出最相似的 k 个。
它是许多研究和工业领域的一个问题，如三维物体渲染，基于内容的图像检索，统计（熵和分歧估计），生物学（基因分类）等。
一般的解决方法是串行的算出相似度，然后用容量为 k 的最小堆（或最大堆）来找找出最相似的 k 个。但这种传统的计算方法往往耗时很长。让人们难以接受。
而每两个向量之间的相似度计算是独立的，这就让这个问题可以很好地被并行的实现。

---

2. GPU 并行计算 kNN

2.1 并行计算思路

kNN 搜索主要分为两个部分：
1. 计算相似度。
2. 找到最相似的 k 个。
因为相似度的计算是独立的，所以这一步可以很好的被并行计算。可以让一个线程计算一个相似度。
而找到最相似的 k 个，则涉及到了排序，一个 query 的排序计算不能被并行，只能让一个线程去计算一个 query 的排序。
而 kNN-CUDA 的作者在他们的论文 Fast k Nearest Neighbor Search using GPU 中研究发现，在 k 很小的时候应该用插入排序，在 k 很大的时候应该使用梳排序（Comb Sort）。他们实验的结果如下图：

而 kNN-CUDA 中则只实现了插入排序。

2.2 源码分析说明

2.2.1 总体说明

kNN-CUDA 一共实现了四种算法，带序号和不带序号的，其中又有为 C 实现的和为 MATLAB 实现的。
我分析的是带序号的为 C 实现的，即 knn_cuda_with_indexes.cu。
这个程序会计算相似度并返回最近的 k 个点的距离和它们的序号。该程序中的函数有：
（并行函数）
- cuComputeDistanceTexture：计算矩阵 A（reference points）和 B（query points）中向量的相似度，其中矩阵 A 是一个 texture。
- cuComputeDistanceGlobal：功能同上，但是 A 是一个矩阵而不是 texture。
- cuInsertionSort：插入排序的实现。
- cuParallelSqrt：计算前 k 个距离矩阵中的平方根。
（一般函数）
- printErrorMessage：输出错误信息。
- knn：计算 knn。
- mexFunction：给 MATLAB 用的接口，只在 MATLAB_CODE == 1 时编译。
- main：kNN 搜索的例子。

2.2.2 具体函数分析

* cuComputeDistanceTexture *

__global__ void cuComputeDistanceTexture(int wA, float * B, int wB, int pB, int dim, float* AB){
    unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int yIndex = blockIdx.y * blockDim.y + threadIdx.y;
    if ( xIndex<wB && yIndex<wA ){
        float ssd = 0;
        for (int i=0; i<dim; i++){
            float tmp  = tex2D(texA, (float)yIndex, (float)i) - B[ i * pB + xIndex ];
            ssd += tmp * tmp;
        }
        AB[yIndex * pB + xIndex] = ssd;
    }
}

功能：
该函数计算矩阵 A