本篇文章来自于ICDE2022,主要工作是提供了 GPU 加速图索引构建和搜索,相比 CPU 以及 GPU 方案 SONG 性能有了更大的提升,更加的充分的利用了线程块内部线程的并行和线程块之间的并行。
今天分享的内容对上次分享的 GPU 加速进行一些补充,还没看过上篇文章的小伙伴可以点击下方链接先看看。
研究背景
SONG 虽然提出了 GPU 在图索引方面构建和搜索的方案,但仍然遵循基于 CPU 的解决方案的搜索范式。例如优先级队列和哈希表的搜索和更新都是可以在 CPU 有效地实现。
然而, GPU 难以动态地维护具有高不规则依赖性的数据结构。虽然已经提出了一组特定于 ANN 的优化技术来消除动态 GPU 内存分配并以更少的 GPU 内存消耗来交换计算,但 SONG 在其实现中依赖于每个线程块中的单个线程进行数据结构操作以获得良好的整体性能。
本文设计了一个 GPU 友好的邻近图搜索算法,使得所有关键步骤都能高效、充分地利用GPU的巨大并行性。其核心思想是使用两种惰性策略:惰性更新和惰性检查,使得数据结构的维护可以并行进行。
GPU 加速图索引构建
图索引的构建依赖于所有待插入节点会排队按顺序插入,从而导致这一步骤对于 GPU 及其不友好。在本篇文章中提出了基于 GPU 的图索引构建方案,在不牺牲图质量的情况下充分利用 GPU 大规模并行的优势。
简单来说,构图阶段可以采用分而治之的思想。将所有点划分为相同大小的不同社群。
在第一阶段中,采用每个线程块处理不同社群。并且每张局部图 Gi 的构造都按照插入点的顺序依次插入。
在第二阶段中,通过一定的合并规则来进行图合并。
在构图过程中一定要确保:
(1)图构造中的所有操作对GPU友好;
(2)实现了块间并行和块内并行;
(3)所得到的图的质量与通过顺序插入构造的NSW图相同;
如下图所示
每个子图的构建则按顺序逐点插入,同时保证每个点的出度最多为 d。构建好子图后,将 t 个图从后向前依次合并。t 个子图的合并则需要t次迭代。每次迭代会做以下操作:
在第 i 次迭代中,会将子图 Gi 中的每个顶点 vij 都分别分配到一个线程块中,因为每个顶点的计算相互独立,所以可以保证线程块间的并行。每个 vij 在根据当前图 G0 检索出 d 个最近邻,用这个点集合和在子图 Gi 中 vij 的邻居集合合并为 vij 的新邻居集合,并且这里面的所有边都会保存在一个边集合 E 中。E 位于 global memory 中用于线程共享。E 的内存开销是 O(n*d)。
为了节省这个邻接矩阵的内存开销,采用稀疏矩阵压缩的存储格式 CSR 来进行存储。
CSR 存储格式如下图所示:
总结:
利用后向边可以并行计算来加速,同时在构建后向边的过程中保存边的信息到global memory的 E 列表中,在为每个点遍历E来寻找前向边。
双调排序
假设我们有双调序列 A,A 序列满足先单调增后单调减(先单调减后单调增)
第一次排序:i 取[0,n/2] 将 i 和 2i+1 比较大小。循环一遍后实质上我们得到了4个单调序列;第二次排序:重复每两组之间的两两比较,循环一遍后我们得到了8个长度为2的单调序列;第三次排序:相邻组之间继续两两比较,得到排序后的数组;
那么我们如何从初始的无序数组转换成双调序列 A 呢?
这过程中需要两次比较,第一次是两组短序列对应位比较大小;第二次是相邻位比较。
初始无序数组 A*,相邻位之间比较,形成一个相邻两组之间单调性相反的序列;继续合并,相邻的两个长度为2的单调序列合并为长度为4的单调序列;重复合并过程,把相邻的两个长度为4的单调序列合并为长度为8的单调序列。
实验
实验选择以下数据集进行图索引的构建和测试。
在不同数据集下,与 SONG 的查询效果做比较,保证相同的召回率时,QPS 最高可达到8倍提升。
在不同的 K 值选择下,GANNS 也表现出了更高的吞吐量。
在图的构建方面,构建时间最多能达到 GPU 的83.5倍,相比 SONG 最高也可达35倍。衡量图质量的方法是相同的查询在 NSW 图和 GANNS 图是否能达到相同的查询,实验结果看到与原图质量几乎一致。
参考文献
[1] Y. Yu, D. Wen, Y. Zhang, L. Qin, W. Zhang and X. Lin, “GPU-accelerated Proximity Graph Approximate Nearest Neighbor Search and Construction,” 2022 IEEE 38th International Conference on Data Engineering (ICDE), Kuala Lumpur, Malaysia, 2022.
[2] Zhao W , Tan S , Li P . SONG: Approximate Nearest Neighbor Search on GPU[C]// 2020 IEEE 36th International Conference on Data Engineering (ICDE). IEEE, 2020.
扫码关注公众号回复【加入社区】进入社群,和大家一起交流讨论。看到这里的你不妨点亮点赞,留下你的足迹吧!