文献分享: BANG——DiskANN在单GPU上的实现与优化

原论文: BANG: Billion-Scale Approximate Nearest Neighbor Search using a Single GPU

---

$\text{PART }$Ⅰ: 导论与预备知识

$\text{1. }$导论

$\text{1.1. }$关于$\text{ANN}$

1️⃣高维$k$最邻近查询

1. 精确查询$\text{(NN)}$：
   • 含义：找到与给定查询点最近的$k$个数据点
   • 困难：由于维度诅咒$\to$难以摆脱暴力扫描$O(n\ast \text{dimension})$的复杂度
2. 近似查询$\text{(ANN)}$：
   • 核心：通过牺牲准确性来换取速度，以减轻维度诅咒
   • $\text{On GPU}$：大规模并行处理可以提高$\text{ANN}$吞吐量(固定时间内的查询数量)
3. 基于图的$\text{ANN}$：
   • 处理大规模数据最为高效的$\text{ANN}$方法
   • $\text{Vamana/DiskANN}$是目前最先进的基于图的$\text{ANN}$(详细的设计$\text{Click Here}$)

$\text{1.2. }$$\text{ANN}$的$\text{GPU}$实现难点

1️⃣$\text{GPU}$的内存有限

1. 含义：目前主流$\text{GPU}$内存有限，无法将构建好的图结构完整载入
2. 现有方案：

   方案	描述	缺陷	文献
   分片	将图分片$\to$不断在$\text{CPU}\leftrightarrows \text{GPU}$交换片以处理整个图	$\text{PCIe}$带宽不够	$\text{GGNN}$
   多$\text{GPU}$	将图有效分割到所有$\text{GPU}$上以容纳并处理整个图	硬件成本高	$\text{SONG}$/$\text{FAISS}$
   压缩	压缩图数据维度使图结构能北方进$\text{GPU}$内存	召回率下降(只适合小数据)	$\text{GGNN}$

2️⃣最有硬件使用

1. $\text{GPU}\leftrightarrows \text{CPU}$负载平衡：确保二者持续并行工作不空闲，并且数据传输量不超过$\text{PCIe}$极限
2. 主存占用：基于$\text{GPU}$的$\text{ANN}$搜索占用的内存显著增加

$\text{1.3. BANG}$的总体优化思路

1️⃣硬件优化

1. 总线优化：减少$\text{CPU-GPU}$间$\text{PCIe}$的通信量$\to$提高吞吐

   优化思路	具体措施
   减少(总共的)总线传输次数	负载平衡，预取/流水线(让$\text{CPU/GPU}$尽量没空闲时间)
   降低(一次的)总线传输量	传输$\text{PQ}$压缩后的向量(而非原始向量)

2. $\text{GPU}$内存优化：避免存放图结构$+$只存放$\text{PQ}$压缩后的向量

2️⃣计算优化

1. 加速遍历/搜索：使用$\text{Bloom}$过滤器，快速判断$a\text{∈}A$式命题的真伪
2. 加速距离计算：使用$\text{PQ}$压缩后的向量计算距离

3️⃣软件优化：设立微内核，将距离计算/排序/更新表操作拆分成更原子化的操作，以提高并行化

$\text{2. GPU}$架构与$\text{CUDA}$编程模型

$\text{2.1. }\text{GPU}$体系结构

1️⃣计算单元组织架构

结构	功能
$\text{CUDA}$核心	类似$\text{ALU}$(但远没$\text{CPU}$的灵活)，可执行浮点运算/张量运算/光线追踪(高级核心)
$\text{Warp}$	多核心共用一个取指/译码器，按$\text{SIMT}$工作(所有线程指令相同/数据可不同)
$\text{SM}$	包含多组$\text{Warps}$，所有$\text{CUDA}$核心共用一套执行上下文(缓存)$\&$共享内存

2️⃣存储层次架构：

1. 不同$\text{SM}$能够$\text{Access}$相同的$\text{L2 Cache}$
2. 显存与缓存之间的带宽极高，但是相比$\text{GPU}$的运算能力仍然有瓶颈

$\text{2.2. }$$\text{CUDA}$编程模型

1️⃣$\text{CUDA}$程序简述

1. $\text{CUDA}$程序的两部分

   程序	运行位置	主要职责
   Host程序	$\text{CPU}$	任务管理/数据传输/启动$\text{GPU}$内核
   Device程序	$\text{GPU}$	执行内核/处理数据

2. $\text{Kernel}$即在$\text{GPU}$上运行的函数，如下简单内核定义示例

   //通过__global__关键字声名内核函数
   __global__ void VecAdd(float* A, float* B, float* C)
   {
      int i = threadIdx.x;
      C[i] = A[i] + B[i];
   }
   int main()
   {
      //通过<<<...>>>中参数指定执行kernel的CUDA thread数量
      VecAdd<<<1, N>>>(A, B, C); 
   }
   

2️⃣线程并行执行架构

1. 线程层次：

   结构	地位	功能
   $\text{Thread}$	并行执行最小单元	执行$\text{Kernel}$的一段代码
   $\text{Warp(32Threads)}$	线程调度的基本单位	所有线程以$\text{SIMD}$方式执行相同指令
   $\text{Block}$	$\text{GPU}$执行线程基本单位	使块内线程内存共享/指令同步
   $\text{Grid}$	并行执行的最大单元	执行整个内核(启动内核时必启动整个$\text{Grid}$)

2. 线程在计算单元的映射：线程层次$\stackrel{层次对应}{\leftrightarrow }\text{GPU}$物理架构
   
   • 注意$\text{SM}$和$\text{Block}$不必$\text{1v1}$对应也可$\text{Nv1}$对应
3. 线程在存储单元的映射

   线程结构	可$\text{Access}$的内存结构	访问速度
   $\text{Thread}$	每线程唯一的$\text{Local Memory}$	极快
   $\text{Block}$	每块唯一的$\text{Shared Memory}$(块中每个线程都可访问)	较快
   所有线程	唯一且共享的$\text{Global Memory}$	较慢

$\text{2.3. CPU}$与$\text{GPU}$

1️⃣$\text{CPU/}\text{GPU}$结构对比

	$\text{GPU}$	$\text{CPU}$
$\text{ALU}$	功能强但数量少(只占$\text{GPU}$小部)，时钟频率极高	功能弱但数量大，时钟频率低
$\text{Cache}$	容量大并分级，缓存后续访问数据	容量很小，用于提高线程服务
控制	复杂串行逻辑，如流水/分支预测/乱序执行	简单(但大规模)并行逻辑

3️⃣$\text{CPU}\underset{数据/指令传输}{\overset{\text{PCIe}}{\leftrightarrow }}\text{GPU}$交互

设备	逻辑地位	$\text{IO}$模块	任务分配
$\text{GPU}$	外设	$\text{IO Block}$(南桥)	控制逻辑和任务调度
$\text{CPU}$	主机	$\text{Copy Engine}$	执行大量并行计算任务

---

$\text{PART }$Ⅱ: $\text{BANG}$的设计

$\text{1. BANG}$的总体设计

$\text{1.1. BANG}$的索引架构

$\text{1.1.1. }$$\text{BANG}$索引(硬件)布局

结构	功能
$\text{RAM}$	存放$\text{Vamana}$算法构建的图结构$+$数据点
$\text{GPU}$内存	存放$\text{Vamana}$算法构建的图中点经过$\text{PQ}$压缩后的向量
$\text{CPU-GPU}$总线	传输压缩向量$\&$协调并行

$\text{1.1.2. BANG}$索引构建算法: $\text{Vamana}$图

1️⃣$\text{Vamana}$图构建基本操作

1. 图查询算法：贪心搜索$\text{GreedySearch}\left(s,{\mathrm{x}}_q,k,L\right)$
   
2. 图剪枝算法：健壮性剪枝$\text{RobustPrune}(p,R,\alpha ,L)$
   

2️⃣$\text{Vamana}$图构建总体流程

$\text{1.1.3. BANG}$索引构建方法: 类似$\text{DiskANN}$架构

1️⃣构建步骤：面向面向内存空间的优化

1. 划分：用$k\text{-means}$将$P$分为多簇(每簇有一中心)，再将$P$所有点分给$\ell \text{>}1$个中心以构成重叠簇
2. 索引：在每个重叠簇中执行$\text{Vamana}$算法，构建相应有向边
3. 合并：将所有构建的有向边合并在一个图中，完成构建

2️⃣关于重叠分簇：为了保证图的连通性，以及后续搜索的$\text{Navigable}$

$\text{1.2. BANG}$的查询架构

$\text{1.2.1. }$第一阶段: 初始化$\text{\&PQ}$表的构建

1️⃣执行的操作

1. 并行化：为查询集$Q_{\rho }$中的每个查询 { q 1 , q 2 , . . . , q ρ } \{q_1,q_2,...,q_{\rho}\} { q1​,q2​,...,qρ​}分配一个独立的$\text{CUDA}$线程$\text{Block}$
2. 距离表：在每个线程块上为每个$q_i$计算并构建$\text{PQ}$距离子表，最终合并$\rho$个子表为距离表
3. 搜索起点：每个$q_i$从图质心开始，即$\text{CPU}\stackrel{传输给}{\leftarrow }\text{WorkList}\stackrel{放入}{\leftarrow }{\text{u}}_i^{\ast }\text{(当前/候选点)}\stackrel{初始化}{\leftarrow }\text{Centroid}$

2️⃣$\text{PQ}$表构建的时序逻辑

时期	操作
查询开始前	将查询点送入$\text{GPU}$的$\text{Copy}$引擎，在$\text{CUDA}$核心上计算/构建/存储距离表
查询开始后	保留距离表在$\text{GPU}$上直到查询结束

$\text{1.2.2. }$第二阶段: 并行$\text{GreedySearch}$主循环

1️⃣前$\text{CPU}$阶段：$\text{CPU}$从内存中获取当前在处理节点$u_i^{\ast }$的邻居集$N_i$

🔁数据传输：$\text{CPU}\stackrel{邻居集N_i}{\to }\text{GPU}$

2️⃣中$\text{GPU}$阶段：接收$u_i^{\ast }$的邻居集$N_i$后，并行地执行内核$\text{\&}$全精度向量的异步传输

1. 执行内核：按顺序执行以下内核及操作

   步骤	操作	内核与否
   过滤邻居	用$\text{Bloom}$并行检查$\forall n\text{∈}{N}_i$中未被访问点$\to$并放入$N_i^{\prime }$(未访问集)$\text{+}$更新$\text{Bloom}$	✔️
   距离计算	用$\text{PQ}$距离表并行计算所有未处理邻居$n_k\text{∈}{N}_i^{\prime }$与查询点$q_i$距离，并存在${\mathcal{D}}_i[k]$	✔️
   邻居排序	将$N_i^{\prime }$和${\mathcal{D}}_i[k]$按与$q_i$的距离执行归并排序，得到排序后的距离${\mathcal{D}}_i^{\prime }$和节点${\mathcal{N}}_i^{\prime }$	✔️
   合并列表	合并当前$\text{WorkLisk}({\mathcal{L}}_i)$与新排序的节点列表${\mathcal{N}}_i^{\prime }$形成新的${\mathcal{L}}_i$	✔️
   更新节点	又将${\mathcal{L}}_i$排序后选取最近的未访问点$u_i^{\ast }$作为下一个当前节点	❌

2. 异步传输：执行内核的同时，$\text{CPU}$将$u_i^{\ast }$的全精度向量传输给$\text{GPU}$$\to$以便后续重排

🔁数据传输：$\text{CPU}\stackrel{当前节点u_i^{\ast }}{\leftarrow }\text{GPU}$

3️⃣后$\text{CPU}$阶段：若${\mathcal{L}}_i$中所有点都被访问过且$|{\mathcal{L}}_i|\text{=}t$，则认为已经收敛$\to$结束循环

$\text{1.2.3. }$第三阶段: (搜索收敛后的)重排与输出

1️⃣重排与输出

1. 重排的时序逻辑

   时间	操作	位置
   搜索过程中	用一个数据结构，存储每个$\text{Iter}$中向$\text{CPU}$发送的全精度候选点	$\text{CPU→GPU}$
   搜索完成后	计算所有候选点到查询点距离，按全精度距离排序后选取前若干	$\text{GPU}$

2. 输出：选取重排后的${\mathcal{L}}_i$中，离$q_i$最近的$k$个节点$\to$作为$k\text{-}$最邻近返回

2️⃣重排的意义：用小成本(仅极小部分即候选点以全精度送往$\text{GPU}$)，补偿由压缩距离产生的误差

$\text{2. BANG}$的微内核设计与并行优化

$\text{2.0. }$微内核总体设计概览

1️⃣设立独立微内核的操作：

阶段	有独立微内核的操作
第一阶段(建表)	$\text{PQ}$表构建操作
第二阶段(主查询)	过滤邻居，距离计算，邻居(归并)排序，归并列表
第三阶段(重排)	重排操作

2️⃣动态线程块的优化：

1. 每个查询分配到一线程块执行，查询过程会依次执行多个内核
2. 执行不同内核时按经验调整线程块大小(如计算密集型内核的块更大)，以保证$\text{GPU}$的高占有

$\text{2.1. }$第一阶段: PQ \textbf{PQ}