knn 搜索的终极加速方法,从此算力不再虚标!

knn 搜索的终极加速方法,从此算力不再虚标!

作者:hws000(hws.000#163.com)
声明:版权所有,转载请联系作者。
出处:https://blog.simbot.net/2023/03/03/knn_search_acceleration_method/

摘要

        本文创造性的提出一种新的knn(k nearest neighbor)搜索加速算法,可以让硬件的利用率达到90%以上。本算法将knn搜索从计算密集、内存IO密集算法,改进为纯计算密集算法,充分挖掘硬件能力,只要堆更多的硬件,就能实现更快的计算。新算法不但可以运行在GPU上,还可以高效运行于各种专业矩阵加速卡上。

1.引言

        knn搜索是精确搜索,通常作为复杂算法的基础,如knn图构建[6]、聚类算法(如k-means)等。这些算法常用于人脸识别[2]、图像搜索[3]、数据挖掘等场景。knn搜索中有两种常用的距离计算方式:内积、欧式距离。在批量计算时,查询数据表示为M,数据库数据表示为N,数据维度表示为K,内积距离计算方法即MN,此公式与矩阵乘等价;欧式距离的计算方法为Mnorm+Nnorm-2MN,可以变换为2(Mnorm/2+Nnorm/2-MN),Mnorm/2和Nnorm/2可以提前计算并作为常量数据读取,此公式等同于矩阵乘和三元加。因此两种距离计算方式均可通过矩阵乘法加速。knn搜索数据集的维度通常介于96~512之间,很少有数据集能超过512的维度。Jia Zhe等人分析了矩阵乘法在NVidia Turing T4 GPU上的执行效率[1],在K比较低时,随着K值增大计算效率逐渐升高。硬件的计算能力是固定的,为什么不同的K值会表现出不同的计算效率?主要问题出在结果的输出上,相同的计算量(即2MNK),不同的K会产生不同的结果个数,内存带宽瓶颈导致K值比较低时,计算结果不能及时输出。knn搜索不仅要计算距离,还要对距离进行排序,才能最终得到TopK个结果,而排序必定需要再次读取计算结果,使内存带宽的瓶颈被进一步放大。Jeff Johnson等人将GPU用于向量搜索[6]时,GPU算力远低于现在,内存带宽的问题并不明显,但TensorCore出现以后,GPU算力直接提升数倍,而内存带宽的提升却有限。我们提出一种计算距离并立即过滤的复合函数,在K值较低时也可以有很高的计算效率,并使用很低的内存带宽,而且能对过滤结果的数量进行控制,从而提高排序性能。

本文作出了以下贡献:
提出一种计算距离并立即过滤的复合函数,只输出符合过滤条件的结果,减少对内存带宽的依赖;
设计一套高效的更新过滤阈值的方法,保证过滤结果准确、可控,同时减少排序的计算量;
针对维度较低的数据,提出一种让查询数据常驻寄存器的方法,进一步提升硬件利用率。

2.研究方法

算力与内存带宽

        在进行内积距离的knn搜索时,假设每次处理的查询量为M条,数据库的数据量为N条,数据维度为K,距离的计算量为2MNK FLOPS,输入、输出类型均为half(FP16 占用2字节),结果输出为2MN字节,输入数据为2(MK+NK)字节。当M远大于K,并远小于N时,输入数据远小于输出数据,内存读写量近似为2MN字节。即当计算量为2MNK FLOPS 时,内存读写量为2MN字节。RTX 2080ti的内存带宽为0.6TB/s,当K为128时,只需76.8TFLOPS的算力就可以产生0.6TB的结果(0.6×128=76.8),而2080ti的理论算力是107.6,此时的硬件利用率为71.4%。数据库的数据规模N一般都在百万以上,knn搜索时TopK的取值一般在1~128之间,使得TopK远小于N。因为大多数距离计算结果不会出现在最终的搜索结果里,所以对距离计算结果进行过滤,可以大大减小输出的字节数。

阈值的产生与更新

        假设数据是随机分布的,搜索结果也会近似均匀的分散在N中。可以将N分为多段,逐段搜索,并使每个分段大小均等于已经搜索的数据量。当N为131072,TopK为128时,分段大小n可以设为以下值:

分段 0 1 2 3 4 5 6 7 8
n 512 512 1024 2048 4096 8192 16384 32768 65536

表1

        处理分段0时,没有过滤阈值,对n条数据进行距离计算、排序,保存前TopK个结果,将最后一个结果的距离值设置为阈值T。处理分段1~S时,使用阈值T对距离计算结果进行过滤,将新过滤的结果与旧的结果合并、排序,用第TopK个结果的距离值更新阈值T,重复此步骤,直到所有数据处理完毕。由于搜索结果分布均匀,前n次过滤得到TopK个结果,再过滤n次,结果个数也应该跟TopK接近。因为过滤的结果数不会完全等于TopK,可能会大于,也可能会小于,并且需要保留上次排序的结果,所以结果的存储空间应该大于2TopK,如设置为4TopK。

        由于每个分段均可过滤出TopK个结果,依次处理完所有分段才能得到最终结果,造成排序负担较重。在实际使用时,可以在初期使用较少的过滤目标数,并逐步增加到TopK。当TopK为512时,过滤目标数、获取阈值T的位置等参数可以使用表2中的值。其中分段4、6、8、10与分段3、5、7、9的目标数相同,属于矫正分段,是为了防止上一个分段没有过滤出目标数量的结果,导致过滤阈值T持续不更新。为了获得更快的排序效果,可以进一步减少矫正分段的比例。

分段 0 1 2 3 4 5 6 7 8 9 10
n 128 128 256 512 1024 2048 4096 8192 16384 32768 65536
目标数 32 32 64 64 128 128 256 256 512 512
T位置 32 32 64 64 128 128 256 256 512 512

表2

需要注意的是:为确保数据随机分布,可能要预先随机调换每条数据的存放位置,否则会导致过滤出的结果数不可控。

查询数据常驻寄存器

        在CUTLASS介绍[8]和mma指令文档[11]中,详细说明了TensorCore的使用方式。一条mma指令需要整个warp中的32个thread协作完成,每4个thread成为一组,每组thread的数据组合为一个m或n,共分为8组,使每条指令一次至少处理8个m或n。输入、输出的数据类型均为half时,每个寄存器存储的数据个数nk=2,4个thread总共存储的k=8。一条mma指令一次处理8个m和n时,产生8*8=64个结果,每个thread存储的结果个数为2,使用的寄存器个数为1。

        当一个warp完成M=64,N=64,K=128的乘法时,通用方法是将步长k设置为8,逐步加载M、N,计算乘积,并保存M*N个结果。当一条mma指令一次只能处理8个m和n时,需要对M、N进行切块,M和N的切块个数分别为TileM=M/8=8,TileN=N/8=8,每个thread使用的寄存器个数分别为A=TileM*1=8、B=TileN*1=8、C=TileM*TileN*1=64。存储数据用的寄存器总数为8+8+64=80。这是一个通用方法的简略描述,实际使用时会用更复杂的操作步骤,来提高性能,具体的内容可以在cuda-samples的cudaTensorCoreGemm等示例中查看。

        在4个thread中存储完整的m、n时,每个thread要使用的寄存器个数为K/nk/4=16。当M有多个切块时,使用的寄存器个数为K/nk/4*TileM=128。N可以使用最小值8逐步加载,使用的寄存器个数为K/nk/4=16,存储结果的寄存器个数 C=TileM*1*1=8。存储数据用的寄存器总数为128+16+8=152。通过调节TileM的大小,可以适用不同的较小的K值,从而使查询数据M常驻寄存器。

使查询数据M常驻寄存器有如下优点:
使数据加载指令数减半,每次只需要加载数据库的数据N即可;
提高各级cache的利用率,cache中只需存储N和结果相关的数据;
不需要同步操作,可以让不同warp的计算指令与过滤指令并发执行,现有的通用方法在加载M、N数据时,需要进行同步,才能有比较好的加载效率。

3.结果与讨论

        RTX 2080ti是Nvidia Turning架构的GPU,RTX 3060是Nvidia Ampere架构的GPU,这两款GPU均可以使用TensorCore进行矩阵计算加速。cuBLAS是Nvidia提供的加速矩阵计算的函数库,包含cublasHgemm、cublasSgemmEx、cublasGemmEx等函数,这些函数的输入、输出类型,以及在RTX 2080ti、RTX 3060上的理论算力(单位:TFLOPS)如下表:

函数名称 输入 输出 RTX 2080ti 算力 RTX 3060 算力
cublasSgemmEx half float 53.8 25.6
cublasHgemm half half 107.6 51.2
cublasGemmEx int8 int 215.2 102.4

表3

        Faiss是Facebook开源的用于高效相似性搜索和密集向量聚类的库,我们测试的版本是v1.7.1。其中的Flat算法使用cublasSgemmEx函数进行矩阵计算加速,并通过专门优化过的排序函数排序。

        我们测试了以上函数,以及我们实现的knn搜索函数(输入、输出与cublasHgemm、cublasSgemmEx相同),在K值为128~1024时的效率。距离计算方式为内积时,算力=2MNK/耗时,其中knn搜索的耗时包含排序的时间。knn搜索默认使用的TopK为16,Faiss的Flat算法使用的TopK为1,并且专门测试了TopK为1~1024时两个算法的效率。我们的knn搜索算法使用与表2相同的方式增加过滤目标数、更新阈值。

        我们使用的排序算法是双调排序,代码参考cuda-samples的bitonicSort算法。我们对排序算法做了2处改进:排序过程实际是过滤出的新数据和旧数据的合并过程,旧数据是排序过的数据,所以不需要重复排序;当过滤出的数据多于目标数据时,参与排序的数据量大于要求的数量,此时只需保证要求数量的排序结果准确,超出部分不需进行完整的排序。

        Deep1B[5]是big ann benchmarks使用的数据集之一,我们同样使用它作为主数据集,它的默认维度是96,数据类型float。我们将它作为各种维度的数据来源,如连续读取128、256个float数据,转换为half类型或者量化为int8类型,然后作为128、256维度的数据使用。这样处理会导致这些数据失去本身的含义,但这样做对单纯的性能测试没有影响。我们同样测试了其它的一些数据集,如gist、glove等。

Figure 1:通过以上测试结果可以看到,在K值较低时我们的算法也有很高的硬件利用率,当查询数据可以常驻寄存器时利用率更高。

        使用Nsight Compute工具对计算与过滤函数进行性能分析,2080ti的硬件利用率最高为91.53%,此时TensorCore利用率为86.94%,3060的硬件利用率与TensorCore利用率相同,最高为96%。在分析工具的Pipe Shared Cycles Active的提示中可以看到,2080ti的fp16与tensor共享一个pipe,导致过滤指令和计算指令产生竞争,这可能是TensorCore的利用率低于硬件利用率的原因。由于欧氏距离需要对矩阵乘法的结果做进一步的计算,才能转换为距离,因此编译出的程序和内积距离有较大差异,这些差异导致部分维度时欧式距离的计算效率更高。

        在对M个查询数据进行搜索的过程中,每个分段均会产生M组过滤结果,每组结果的个数为x,M个x组合成集合X,x的平均值近似等于过滤目标数。对集合X进行分析后得到如下结果:随机变量x服从自由度为tk的卡方分布,其中tk近似等于过滤目标数,与数据维度等参数无关。卡方分布是一种特殊的Gamma分布,此时Gamma分布的形状参数α=tk/2、尺度参数β=2。

Figure 2:上部截图为搜索算法的一张执行截图,数据库大小209万,查询量21504,数据维度96,数据类型half,TopK为128,循环16次结束。第1、2列是距离计算与过滤函数的执行时间(单位:ms)和相应算力(单位:TFLOPS);第3、4、5列分别是结果解析函数、排序函数、阈值更新函数的执行时间;第6列是分段大小;第7、8列是过滤结果数的均值和标准差;第9、10列是由均值、标准差计算出的Gamma分布的形状参数α、尺度参数β;第11、12列是过滤结果数的最大值、最小值;第13列是过滤目标数。最后一行是整个查询的耗时和相应算力。

Figure 3:左侧蓝色实线是tk=32的概率密度曲线,蓝色圆点是目标数为32的过滤结果数的统计值;橙色实线是tk=64的概率密度曲线,橙色圆点是目标数为64的过滤结果数的统计值;绿色实线是tk=128的概率密度曲线,绿色圆点是目标数为128的过滤结果数的统计值。

        由于首个分段是完整的输出、排序,并存储到结果队列,因此结果队列永不为空,所以初始过滤目标数范围内的召回率为1。当使用表2的方式增加过滤目标数时,由于多个矫正分段的存在,使TopK的召回率也接近于1。可以通过结果队列的填充情况推算召回率,当填充位置达到或超过TopK时,召回率为1。通过控制过滤的结果数,可以减小输出带宽,同时也减小了排序的次数和每次排序的计算量,使硬件的性能得到充分利用。

        Felix Chern等人在Google TPU上实现了峰值性能的knn搜索方法[7],核心方法也是对结果进行过滤,不过他们没有使用迭代更新阈值的策略,而是通过召回率来控制过滤结果的数量,消除了迭代、排序的开销。他们对内存瓶颈、指令瓶颈做了很详细的分析。

4.结论

        虽然GPU有很高的理论算力,但很少有应用能将其充分利用起来。在本文中,我们提出了一套以过滤为基础的算法,通过对输出结果进行过滤,减少输出带宽,提高硬件利用率;通过对过滤结果数量的控制,提高排序速度,但几乎不影响召回率;通过专门针对低维度计算的优化,使新硬件的利用率达到96%。

本文的部分内容来自我的blog[9]和专利。
Ding Nan[4]、英伟达[10]等对GPU的内存瓶颈、指令瓶颈也有很详细的分析。

参考文献

[1] Jia, Zhe, et al. “Dissecting the NVidia Turing T4 GPU via microbenchmarking.” arXiv preprint arXiv:1903.07486 (2019).

[2] Ebrahimpour, Hossein, and Abbas Kouzani. “Face recognition using bagging KNN.” International conference on signal processing and communication systems (ICSPCS’2007) australia, gold coast. sn, 2007.

[3] Gordo, Albert, et al. “Deep image retrieval: Learning global representations for image search.” European conference on computer vision. Springer, Cham, 2016.

[4] Ding, Nan, and Samuel Williams. “An instruction roofline model for gpus.” 2019 IEEE/ACM Performance Modeling, Benchmarking and Simulation of High Performance Computer Systems (PMBS) (2019): 7-18.

[5] Babenko, Artem, and Victor Lempitsky. “Efficient indexing of billion-scale datasets of deep descriptors.” Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2016.

[6] Johnson, Jeff, Matthijs Douze, and Hervé Jégou. “Billion-scale similarity search with gpus.” IEEE Transactions on Big Data 7.3 (2019): 535-547.

[7] Chern, Felix, et al. “Tpu-knn: K nearest neighbor search at peak flop/s.” arXiv preprint arXiv:2206.14286 (2022).

[8] Kerr, Andrew, et al. “Cutlass: Fast linear algebra in cuda c++.” NVIDIA Developer Blog (2017).

[9] hws000. “The fastest knn search algorithm” https://blog.simbot.net/2022/03/12/the-fastest-knn-search-algorithm/

[10] “Matrix Multiplication Background User’s Guide” NVIDIA Documentation Center https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html

[11] “Matrix multiply-accumulate operation using mma instruction” NVIDIA Documentation Center https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma