向量检索压缩算法 SIMD 万字详解论文阅读 Product Quantization cpu cache

FastScan快速PQ扫描论文阅读笔记 Cache locality is not enough high-performance nearest neighbor search with product quantization fast scan

Ding Ning2025-04-232025-04-23

论文简介

Cache locality is not enough high-performance nearest neighbor search with product quantization fast scan 是一篇发表于 VLDB2015 的经典论文。

主要贡献

1、提出了一个通过寄存器 LUTs 代替 mem LUTs 来加速 PQ 查表的设计，并通过 SIMD 加速优化

论文内容

1. 相关工作：PQ 及 IVFPQ

1.1. Product Quantization

量化器的构建思路：

PQ 的核心压缩思想是让距离向量 x 最近的数据集特征向量 c 代替 x 。如下图：

alt text

q 是一个映射函数，它将一个实数域的d 维向量映射到一个候选集 C 中，对于任何一个输入向量 x ，它都能跟给出一个对应的映射 label ，也就是某个质心 ci 代表 x 。

有了上面的形式化定义，我们还需要思考这个函数具体怎么设计，因此有最简单的映射函数为：

alt text

这里 ||x - ci|| 指的是向量 x 与 ci 之间的欧几里得范数，也就是欧氏距离。q(x) 也就是一个让距离最小的函数。

alt text

因此，一个给定的向量 x ，都可以被编码为距离它最近的数据集质心，也就是编码成集合 C 中的 label （上图中举的例子）。

但是这种粗暴的方法有很大的问题：

1、如果选取的质心数量太小，也就是 |C| 太小，那么大量的向量都被映射到同一个质心上，误差自然很大。

2、如果选取的质心数量太大，虽然很精确，但是 码本大小上升，会导致单个 label 的存储代价增大（如上图的 32 bit 码本，在质心数量为 2^64 时，需要的存储代价增加了一倍）；另外过大的质心数量将导致聚类算法的代价过大，无法实际地应用。

为了解决这个问题， Product Quantization 将整个向量 x 拆分成了若干段，对每段进行独立的压缩，如下图：

alt text

而具体的形式化描述如下图：

alt text

一般来说，子段向量长度取 8 ，质心数量取 2^8=256 能够拿到最好的 质量和复杂度之间的 tradeoff 。

alt text

一个简单的例子：

alt text

这个图代表了database vector 在 8*8 量化后的编码结果。其中 pqr…u 是若干个 database vector 的量化后表示，每 8bits 代表一个小 segment 的编码，每个向量有 8 段。

1.2. ANN Search with Product Quantization

知道了怎么压缩，接下来要讲解的就是怎么计算相似度了。

首先是一个对相似度计算的形式化的描述：

alt text

然后，论文讲解了 Asymmetric Distance Computation (ADC) 计算的具体定义，清晰地理解 ADC 是理解查询流程的必要条件：

alt text

如上图， query vector 是 y ，数据库 vector 是 p 。因此两者的近似距离就是每个小段的距离的和，也就是 y 的 j 小段与 p 的 j 小段的质心之间的距离。

y 与每个质心的距离只需要计算一次，就可以被同一个质心的所有向量 p 共享。

因此，在 IVFADC 索引（倒排 + ADC 相似度计算）中，具体的查询流程为：

alt text

也就是：

1、根据查询向量 y 与分区中心点的距离，选择一个需要扫描的倒排分区，这个分区中很可能有目标的结果向量。

2、计算 LUTs ，在本文中名为 compute distance tables(y)

3、PQ Scan ，扫描整个分区，计算分区中所有向量与查询向量 y 的相似度

LUTs 的定义如下：

alt text

我们举个例子，假设量化参数中每个 segment 的聚类中心数量为 2^8 ，共 8 个 segment ，那么 LUTs 就是一个如下图的 floats 数组：

alt text

其中 00-ff 共 256 个列， D0-D7 共 8 行，每个元素都是一个 float ，代表距离的计算结果。因此 LUTs 是一个 8*2^8 大小的二维float数组。

这里不得不提一嘴，论文中的距离都是用 int 形式来绘图的，也就是上图中左侧的距离 LOOK UP TABLE 中没有一个是浮点数，都是正整数。在论文的实际代码中，这里全都是 float32，详情可以看我 2.3 附录-四种实现 的描述。但是事实上，对于平方欧氏距离，结果必定是正整数，在不会出现数据越界的前提下，也可以使用 unsigned int ，需要根据实际情况取舍。（浮点数的运算复杂度远大于 int ，int 运算一般一个时钟周期即可完成，而 float 运算代价相对较大）

由此，查询向量 y 到 database 中任何一个 vector 的距离都能被查表得出：

alt text

Dj 代表了查询向量 y 在第 j 个 segment 上，与任何一个编码代表的聚类中心点的距离。因此 d(p,y) 就可以被表示为数据向量 p 的量化编码与对应的 y 的距离之和，也就是上面的第二个公式。

这个查表过程可以用示例图的右侧图来理解， y 与第一条向量的距离可以用每个 segment 的距离之和表示。第一个 segment 编码为 02 ，所以去 D0 中找 02 的 float ；第二个 segment 编码为 04 ，所以去第二个 segment D1 中找 04 的 float ……

总体过程是下图的伪代码：

alt text

其中 NNS 是整个查询流程，可以看到还是很简单的。

此外，关于 k* 和 k, m 三个参数之间的关系：

alt text

2. PQ SCAN LIMITATIONS

2.1. Memory Accesses

论文把主要操作符归类为三种：

alt text

分别对应着代码中的绿、紫、粉：

alt text

值得注意的是，操作 1 几乎能够百分百命中 L1 Cache ，原因是 硬件prefetch 能够识别固定模式访问的数据，而 p[j] 的访问是有固定流程 的。

如： p[0] p[1] 之间间隔了固定的指令 d <- balabala ，因此硬件prefetch 能够识别到这种固定模式的访存，可以提前在 d <- balabala 快结束时提前将 p[] 加载至 L1 cache 。

而对于指令 2 ， CPU 访问 cache 的等级会因为 PQ 的压缩配置而有所不同。对于元素数量较小的 LUTs ， L1 cache 能够完全容纳，因此将数据放置于 L1 cache 会最大化加速访存；而对于较大的 LUTs ，只能放于 L3 才能容纳。

论文给出了具体编码和能容纳的对应 cache 级别表：

alt text

我们拿 8 * 8 举例， 8 * 8 的 LUTs 元素数量为 8 * (2^8) ，假设每个元素为 float32 （向量运算中最常用的距离），需要的存储空间为：

(分段数 * 每段质心数) * 距离存储代价
(8 * (2^8)) * 32 bits = 65536 bits = 8 KiB < 32 KiB

可以看到，这是能够装进 L1 的。但如果再增加码本大小，也就是降低 segment 数量，同时增加 segment 的质心数量，比如 4 * 16 ，LUTs 所需存储空间就变成了：

4 * (2^16) * 32 bits = 1024 KiB = 1 MiB

只能放到 L3 cache 中了。因此论文选用了权衡精度和复杂度后，总体最优秀的 8 * 8 来实现。

为了分析扫描算法的瓶颈，论文使用 performance counters 对程序进行分析：

alt text

最重要的是 cycles 和 cycles w/load ，前者是完成指令所需的总时钟周期数量。而后者是在执行特定算法或程序期间，存在未完成的加载操作所需的 CPU 周期数被叫做 pending load operations。具体来说，是指在 PQ Scan（Product Quantization扫描）的不同实现过程中，当有数据需要从内存加载到缓存或者处理器寄存器时，如果这些加载操作未能立即完成（例如由于缓存缺失导致），那么处理器就需要等待这些加载操作完成，这段时间是以CPU周期数来衡量的。

图中 naive, libpq, avx, gather 是四种微观实现方式， naive 是传统的无优化实现，而 libpq 使用 load 一次地址到寄存器然后依次执行移位操作来避免多次 load 指令，此外 avx 和 gather 是两种 SIMD 实现，我们后续会讲解。

可以看到：

alt text

上面花了红框的最重要的扫描时间中，四种实现几乎没有明显的优化效果，尤其是 gather 指令在 SIMD load 操作上的复杂性，反而增加了耗时。

2.2. Inability to Leverage SIMD Instructions

PQ 扫描算法如果想通过 SIMD 优化，一般需要水平展开，也就是在一个计算指令中同时计算多条向量的局部距离。也就是下面这张图：

alt text

一条 c = _mm256_add_epu8(a, b) 指令将同时计算两个 256 位 XMM 寄存器 a, b 中存储的 8 条单位长 8bit 的无符号正整数的对位加和结果，存储到 XMM 寄存器 c 中。（reg-reg模式）

也就是画红框的部分：

alt text

可以看到， SIMD 实现是分 8 路 同步计算 8 条向量的平方欧氏距离 的。其中 :

a , b, c…, h 分别代表 8 条向量
a[0] , a[1] 代表向量 a 的第 0，1 段
D0[a[0]] 代表了查询向量 y 与向量 a 的第一段的平方欧氏距离，这个距离是通过预计算存储在距离 LUTs 中的。

这就出现了一个问题， LUTs 中的 D0[a[0]] 、 D0[b[0]] 往往都是离散的，地址不连续的（a , b … 都是同一个分段 0 中的随机质心 a, b …），如何把这些数据放置于同一个 SIMD 寄存器中呢？

论文提到了两种方法：

直接使用若干个 load 指令，从不同的 memory address 中分别 load 数据：

alt text

使用一个 gather 指令，把 load 操作封装：

上述四个指令图片引自：https://blog.csdn.net/kaka11/article/details/122081261

论文经过测试后，发现这两种实现方式都没什么优化效果， gather 甚至反而造成了性能的下降：

alt text

核心原因还是因为不连续内容的 load 是重量级操作，一口气 load 16 个 item 然后让 CPU 统一执行相同的运算，反而让流水线退化为线性操作，无法有效地利用 CPU 。

2.3 附录-四种基础 Scan 的实现（推荐 SIMD 新手阅读）

fastscan 的官方开源项目中提供了上面四种比较的 scan 函数的实现。我们来阅读下：

1. Naive Scan

PS: 中英文注释都是我加的，直接看我的注释就好了。 Naive Scan 的作用主要是帮我们理清运算流程。

alt text

//
// Copyright (c) 2015 – Thomson Licensing, SAS
//
// The source code form of this open source project is subject to the terms of the
// Clear BSD license.
//
// You can redistribute it and/or modify it under the terms of the Clear BSD
// License (See LICENSE file).
//

#include "scan_naive.hpp"

#include <cfloat>
#include <cstdint>
#include <immintrin.h>

/* 8*8 的 LUTs，共 8 个 segments （NSQ），每个 segment 有 2^8=256 个聚类中心，因此每个 segments 都有 256 个局部距离（NCENT） */
#define NSQ 8
#define NCENT 256				

/**
 * * @brief Scan the partition and find the best candidates using a binheap.
 * * @param partition The partition to scan. 这个 partition 是一个局部 partition，里面包含了很多个 database 向量 ，每个向量用其段编码表示，也就是每个向量都由 NSQ 个 uint8_t 表示。
 * * @param dists The distances between the centroids. 可以理解为一个二维数组 float[NSQ][NCENT]
 * * @param n The number of elements in the partition.
 * * @param pqp The parameters for the PQ.
 * * @param bh The binheap to store the best candidates.
 */
void scan_bh(const char* partition, const float* dists,
		unsigned long n, pq_params pqp, binheap* bh) {
	long binheap_op = 0;
	for (int t = 0; t < bh->capacity(); ++t) {
		bh->push(0, FLT_MAX - t);
	}
	float min = bh->max();
	float candidate;
	uint8_t* u8_buf = (uint8_t*) partition;					// 读取 partition 的时候，直接将其转换为 uint8_t* 类型，这样就可以直接读取每个 segment 的编码了。
	for (unsigned long i = 0; i < n; ++i) {					// partition 中共有 n 个向量，依次遍历，计算距离
		uint8_t* pqcode = u8_buf + i * NSQ;					// 每个向量用 NSQ 个 uint8_t 表示，表示该向量在每个 segment 中的编码
		candidate = 0;
		for (unsigned j = 0; j < NSQ; ++j) {				// 遍历每个 segment，计算该向量在每个 segment 中的距离，加和成为 向量 y 与当前 database vector 的近似距离
			candidate += dists[j * NCENT + pqcode[j]];		// dists[j * NCENT + pqcode[j]] 表示 在第 j 个 segment 中的 pqcode[j]编码 对应的局部距离
		}
		if (candidate < min) {
			bh->push(i, candidate);							// 将当前向量的索引和距离加入到 binheap 中，binheap 会自动排序，得到 KNN
			min = bh->max();
			binheap_op++;
		}
	}
}

2. SSE Scan

//
// Copyright (c) 2015 – Thomson Licensing, SAS
//
// The source code form of this open source project is subject to the terms of the
// Clear BSD license.
//
// You can redistribute it and/or modify it under the terms of the Clear BSD
// License (See LICENSE file).
//


#include <cfloat>
#include <immintrin.h>
#include "scan_sse.hpp"
#include <cassert>
#include <avxintrin.h>

#define NSQ 8
#define NCENT 256
#define SSE_SZ 4

/**
 * * @brief Scan the partition and find the best candidates using a binheap.
 * * @param partition The partition to scan.
 * * @param dists The distances between the centroids.
 * * @param n The number of elements in the partition.
 * * @param pqp The parameters for the PQ. 在这里没有用
 * * @param bh The binheap to store the best candidates.
 * * @details This function uses SSE instructions to speed up the computation of the distances.
 */
void scan_prefetch_sse(const char* partition, const float* dists,
		unsigned long n, pq_params pqp, binheap* bh) {
	long binheap_op = 0;
	for (int t = 0; t < bh->capacity(); ++t) {
		bh->push(0, FLT_MAX - t);
	}
	float minb = bh->max();
	__m128 min = _mm_set1_ps(minb);
	float mcandidates[SSE_SZ];
	__m128 partial;
	uint8_t* u8_buf = (uint8_t*) partition;
	for (unsigned long i = 0; i < n ; i += SSE_SZ) {
		const uint8_t* const pqcode0 = u8_buf + i * NSQ;
		const uint8_t* const pqcode1 = pqcode0 + NSQ;
		const uint8_t* const pqcode2 = pqcode0 + 2*NSQ;
		const uint8_t* const pqcode3 = pqcode0 + 3*NSQ;
		__m128 candidates = _mm_set_ps(
					dists[pqcode3[0]],
					dists[pqcode2[0]],
					dists[pqcode1[0]],
					dists[pqcode0[0]]
					);
			// Such perf critical loop. Pls unroll
			for (unsigned j = 1; j < NSQ; ++j) {
				const float* const cdist = dists + j * NCENT;
				partial = _mm_set_ps(
						cdist[pqcode3[j]],
						cdist[pqcode2[j]],
						cdist[pqcode1[j]],
						cdist[pqcode0[j]]
						);
				candidates = _mm_add_ps(candidates, partial);
			}
			const __m128 cmp = _mm_cmp_ps(candidates, min,_CMP_LT_OS);
			if (!_mm_testz_ps (cmp, cmp)) {
				_mm_store_ps(mcandidates, candidates);
				for(unsigned ii = 0; ii < SSE_SZ; ++ii) {
					if(mcandidates[ii] < minb) {
						bh->push(i+ii, mcandidates[ii]);
						minb = bh->max();
						binheap_op++;
					}
				}
				min = _mm_set1_ps(minb);
			}
		_mm_prefetch(reinterpret_cast<const char *>(pqcode0 + 64*4), _MM_HINT_NTA);
	}
}

#undef SSE_SZ

SSE Scan 使用了 SSE 指令集进行优化， SSE 指令集提供了对 128位 XMM 寄存器 直接操作的汇编指令集，我们将依次讲解这里用到的指令集：

我们可以把 __m128 理解为一个 128位的XMM 寄存器 ，当上下文中同时用到的 __m128 数量不超过 CPU 的 SIMD 寄存器总数量时，不会出现额外的 load&save 开销。

首先是第一个 _mm_set1_ps ：

alt text

可以看到函数注释中已经说清楚了，使用传入的一个 float32 ，设置同一个 128-bit floating-point vector of [4 x float] 的四路浮点数，然后返回这个寄存器的引用。

接着是 load 指令 _mm_set_ps ，它的作用是 从内存中 加载指定的数据至 寄存器：

alt text

在这里加载了 4 个 float32 ，共 128 位数据进入寄存器。

然后是运算指令 _mm_add_ps ：

alt text

它的作用是分四路，分别计算每路中的两个 float32 的 add 结果，将其存放于第三个寄存器中。

这里直接使用 candidates 承接了计算结果，所以计算结果会被直接更新到 candidates 寄存器，而无需额外的寄存器来缓存。

比较操作符 _mm_cmp_ps ，状态设置操作符 _mm_testz_ps：

alt text

比较 + 状态两个指令，共同构成了一个分支判断语句，它的作用是检查这批计算的距离有没有比当前 min 距离更小的距离，如果存在任何一个，就需要取出寄存器中缓存的距离到内存，然后使用 SISD 指令集来读取距离，更新候选结果。

这里额外涉及到一个 EFLAGS寄存器 ， EFLAGS寄存器 是x86架构处理器中的一个重要组成部分，用于存储有关程序执行或运行状态的信息。它是一个32位的寄存器，在较旧的16位架构中称为FLAGS寄存器，而64位架构中则扩展为RFLAGS寄存器（包含EFLAGS的所有功能并增加了更多标志位）。EFLAGS寄存器中的各个位代表不同的状态标志或控制标志，这些标志可以用来影响指令的行为或者提供关于之前指令执行结果的信息。

alt text

存储指令 _mm_store_ps ：

alt text

_mm_store_ps 会把一个寄存器的数据存储到一段线性内存中，一般是在 寄存器数量不足需要内存做缓存 or 需要从 SIMD 模式转换为 SISD 模式指令 情况下使用。

综上，我们可以看到 SSE 优化版的整体思路是，分别从 4 个内存地址 load 数据到一个 4 路 XMM 寄存器，然后每次都是同时更新 4 路的局部距离累加和，最后得到 4 条向量与查询向量的距离。这四个距离如果都比当前的 min 大，那么就都是没用的，直接跳过。否则，还需要缓存到内存中逐个比较距离是否有用，有用则更新距离。

3. AVX Scan

AVX Scan 对比 SSE Scan ，只是拓展了 XMM 寄存器的容量，一次计算 8 个距离而已，我们把代码列在下面，就不赘述了：

//
// Copyright (c) 2015 – Thomson Licensing, SAS
//
// The source code form of this open source project is subject to the terms of the
// Clear BSD license.
//
// You can redistribute it and/or modify it under the terms of the Clear BSD
// License (See LICENSE file).
//


#include <cfloat>
#include <immintrin.h>
#include "scan_avx.hpp"

#define NSQ 8
#define BITSQ 8
#define NCENT 256

#define AVX_SZ 8

/** Base functions **/
__attribute__((always_inline))
static inline void simd_lookup_add_min(const uint8_t* const pqcode0,
		const uint8_t* const& pqcode1,
		const uint8_t* const& pqcode2,
		const uint8_t* const& pqcode3,
		const uint8_t* const& pqcode4,
		const uint8_t* const& pqcode5,
		const uint8_t* const& pqcode6,
		const uint8_t* const& pqcode7,
		const float*& dists,
		__m256& candidates,
		__m256& partial,
		float (&mcandidates)[AVX_SZ],
		__m256& min,
		float& minb,
		binheap*& bh,
		long& binheap_op,
		unsigned long& i
		) {
	candidates = _mm256_set_ps(
			dists[pqcode7[0]],
			dists[pqcode6[0]],
			dists[pqcode5[0]],
			dists[pqcode4[0]],
			dists[pqcode3[0]],
			dists[pqcode2[0]],
			dists[pqcode1[0]],
			dists[pqcode0[0]]
			);
	// Such perf critical loop. Pls unroll
	for (unsigned j = 1; j < NSQ; ++j) {
		const float* const cdist = dists + j * NCENT;
		partial = _mm256_set_ps(
				cdist[pqcode7[j]],
				cdist[pqcode6[j]],
				cdist[pqcode5[j]],
				cdist[pqcode4[j]],
				cdist[pqcode3[j]],
				cdist[pqcode2[j]],
				cdist[pqcode1[j]],
				cdist[pqcode0[j]]
				);
		candidates = _mm256_add_ps(candidates, partial);
	}
	const __m256 cmp = _mm256_cmp_ps(candidates, min,_CMP_LT_OS);
	if (!_mm256_testz_ps (cmp, cmp)) {
		_mm256_store_ps(mcandidates, candidates);
		for(unsigned ii = 0; ii < AVX_SZ; ++ii) {
			if(mcandidates[ii] < minb) {
				bh->push(i+ii, mcandidates[ii]);
				minb = bh->max();
				binheap_op++;
			}
		}
		min = _mm256_set1_ps(minb);
	}
}

void scan_bh_prefetch_avx(const char* partition, const float* dists,
		unsigned long n, pq_params pqp, binheap* bh) {
	long binheap_op = 0;
	for (int t = 0; t < bh->capacity(); ++t) {
		bh->push(0, FLT_MAX - t);
	}
	float minb = bh->max();
	__m256 min = _mm256_set1_ps(minb);
	__m256 candidates = _mm256_setzero_ps();
	float mcandidates[AVX_SZ];
	__m256 partial;
	uint8_t* u8_buf = (uint8_t*) partition;
	for (unsigned long i = 0; i < n ; i += AVX_SZ) {
		const uint8_t* const pqcode0 = u8_buf + i * NSQ;
		const uint8_t* const pqcode1 = pqcode0 + NSQ;
		const uint8_t* const pqcode2 = pqcode0 + 2*NSQ;
		const uint8_t* const pqcode3 = pqcode0 + 3*NSQ;
		const uint8_t* const pqcode4 = pqcode0 + 4*NSQ;
		const uint8_t* const pqcode5 = pqcode0 + 5*NSQ;
		const uint8_t* const pqcode6 = pqcode0 + 6*NSQ;
		const uint8_t* const pqcode7 = pqcode0 + 7*NSQ;
		simd_lookup_add_min(pqcode0, pqcode1, pqcode2, pqcode3, pqcode4, pqcode5, pqcode6, pqcode7,
				dists, candidates, partial, mcandidates, min, minb, bh, binheap_op,
				i);
		_mm_prefetch(pqcode0 + 64*8, _MM_HINT_NTA);
	}
}

4. VGATHER Scan

VGATHER Scan 是论文中提及的 gather 指令，它的作用是在一个指令中同时加载多个地址的数据进入 SIMD 寄存器。其它的步骤几乎和 AVX2 是一致的，我们也不过多赘述了。

//
// Copyright (c) 2015 – Thomson Licensing, SAS
//
// The source code form of this open source project is subject to the terms of the
// Clear BSD license.
//
// You can redistribute it and/or modify it under the terms of the Clear BSD
// License (See LICENSE file).
//


#include <cfloat>
#include <immintrin.h>

#include "scan_gather.hpp"

#define NSQ 8
#define BITSQ 8
#define NCENT 256

#define AVX_SZ 8

void scan_bh_avx_gths(const char* partition, const float* dists, unsigned long n,
		pq_params pqp, binheap* bh) {
	for (int t = 0; t < bh->capacity(); ++t) {
		bh->push(0, FLT_MAX - t);
	}
	float minb = bh->max();
	__m256 min = _mm256_set1_ps(minb);
	float mcandidates[AVX_SZ];
	const __m256i mask_firstbyte = _mm256_set1_epi32(0xFF);
	const __m256i * u256_buf = (const __m256i *) partition;
	// Each iteration of the loop will process a block
	// 1 block = 32 pqcodes
	const unsigned long blk_n = n / 32;
	__m256i data[NSQ];
	// This loop processes a block
	// i: block number
	for (unsigned long i = 0; i < blk_n; i += 1) {
		// Load block data
		for (unsigned int j = 0; j < NSQ; j++) {
			const __m256i* __restrict__ pqcode = u256_buf + i * NSQ + j;
			data[j] = _mm256_load_si256(pqcode);  // Temporal load
		}
		// Prefetch next block. This has no significant impact.
		//_mm_prefetch(u256_buf + (i+1) * NSQ, _MM_HINT_NTA);
		// We handle 8 pqcodes/gather operation
		// 8 pqcodes/gather * 4 = 32 pqcodes / blocks
		for(unsigned s=0; s<4; s++) {
			// Gather and Add
			__m256 candidates = _mm256_i32gather_ps(dists,_mm256_and_si256(data[0],mask_firstbyte),4);
			for(unsigned int j=1;j< NSQ;j++) {
				const float *__restrict__ cdists = dists+j*NCENT;
				const __m256i value = _mm256_and_si256(data[j],mask_firstbyte);
				const __m256 partial = _mm256_i32gather_ps(cdists,value,4);
				candidates = _mm256_add_ps(candidates, partial);
			}
			// Check
			const __m256 cmp = _mm256_cmp_ps(candidates, min,_CMP_LT_OS);
			if (!_mm256_testz_ps (cmp, cmp)) {
				_mm256_store_ps(mcandidates, candidates);
				for(unsigned ii = 0; ii < AVX_SZ; ++ii) {
					if(mcandidates[ii] < minb) {
						const int index=(i*32)+(ii*4)+s;
						bh->push(index, mcandidates[ii]);
						minb = bh->max();
					}
				}
				min = _mm256_set1_ps(minb);
			}
			// Shift to next byte
			for(unsigned int j=0;j<NSQ;j++) {
				data[j] = _mm256_srli_epi32(data[j],8);
			}
		}
	}
}

void gather_check(const char* partition, const char* i32_partition,
		unsigned long n) {
	const __m256i * u256_buf = (const __m256i *) i32_partition;
	const unsigned long blk_n = n / (8 * 32);
	__m256i data[NSQ];
	char* data_mem;
	posix_memalign((void**) &data_mem, 32, 320);
	for (unsigned long i = 0; i < blk_n; ++i) {
		// Load
		for (unsigned int j = 0; j < NSQ; j++) {
			const __m256i* __restrict__ pqcode = u256_buf + i * NSQ + j;
			data[j] = _mm256_load_si256(pqcode);  // Temporal load
		}
		// Check
		for (unsigned int j = 0; j < NSQ; j++) {
			_mm256_store_si256((__m256i*) data_mem, data[j]);
			for(unsigned l=0; l<32; ++l) {
				if(partition[(i*32+l)*8+j] != data_mem[l]) {
					printf("i=%lu, j=%u, l=%u\n", i,j,l);
					exit(32);
				}
			}
		}
	}
	free(data_mem);
}

#undef AVX_SZ

项目地址：https://github.com/technicolor-research/pq-fast-scan

2.4 结论

因此，综合上面的两个小节，论文得出了结论：

1、PQ Scan 是一个 memory access largely 操作。

2、PQ Scan 难以直接应用 SIMD 进行性能优化，虽然使用 SIMD 优化确实让指令的总 Instruction 数量减少了，但反而提升了操作复杂度，难以让调度器通过流水线进行优化，无法充分利用 CPU 的运算性能。

3. PQ FAST SCAN

3.1 论文内容

3.1.1 概述

PQ FAST SCAN 是论文提出的用于代替 PQ SCAN 的上位替换算法。它的整体思路是 先用进一步压缩的近似距离粗粒度过滤，再精细计算 ：

alt text

3.2.2 Vector Grouping

我这里直接引用一张大模型的描述，比我讲得好。 Vector Grouping 的核心是把拥有相似编码表示的向量顺序扫描，这样就不用频繁地访问 LUTs 中的各个部分。这部分 group 操作是在 IVF build 之后，拿到了一个 prtition 的向量，然后进行的。

alt text

简单来说：

输入是向量的 LUTs 编码表示，也就是每条向量都被 8 个 uint8 表示。
输出的是若干个 group 。每个 group 中，向量的前 4 个 segments 的编码都是相同的，每个 group 都是有长度的，表示这个 group 中有多少条向量。输出结果最多有 16^4 个组。

具体每个组的定义如下：

alt text

values 指的就是查询命中的 4 个组的 id 都是什么，比如 values = {0, 2, 3, 14} 就指的是这个 group 中的查询都是在 D0-group0, D1-group2, D2-group3, D3-group14, 中。也就是这个 group 中的全部向量的压缩表示中，第一个 segment 的索引都是位于 0 - 16 ，第二个都位于 2 * 16-3 * 16 , 第三个都位于 3* 16 - 4* 16 ，第四个都位于 14 * 16 - 15 * 16 。

因此，这些查询都只会访问这四个的 segment 的 4 个 group ，每个 group 中都只有 16 条 float32 表示的局部距离。

3.2.3 Minimum Tables

alt text

Minimum Tables 的思路很简单，通过丧失精度压缩 LUTs ，但是这里的丧失精度很有趣，它通过用局部下界代替整体的思路来采样。

具体的实现如上图，红色部分是 LUTs 的每个部分的最小值，用这个最小值来代替这个部分。

在具体实现中，LUTs 中的最小值是 float32 ，论文又采用了 SQ 压缩来将其量化为 int8 ，让每个 D 都能被装进一个 SIMD 寄存器。

3.2 源码阅读

3.2.1 重要实现

论文里的描述我个人认为很难理解，所以我把论文对应的代码额外阅读了，以此来解决论文中的描述问题。实现的 FastScan 为函数 scan_partition_1 ，两个关键步骤都列在了这里：

alt text

文件主要包含了几个函数：

alt text

1. begin_scan

begin_scan 函数是一个全表扫描函数，它会扫描整个 partition 中的部分向量，计算出其中距离的最大值。最大值会用于后续的距离量化。

FORCE_INLINE
static inline float begin_scan(const std::uint8_t* partition, const unsigned* labels, const float* dists,
		unsigned long n, binheap* bh) {
	// 初始化 binheap ，直接加满元素
	for(int pqcode_i = 0; pqcode_i < bh->capacity(); ++pqcode_i) {
		float candidate = 0;
		const std::uint8_t* const pqcode = partition + pqcode_i * NSQ;	// 每次取出一个向量的 8 位编码
		for(int comp_i = 0; comp_i < NSQ; ++comp_i) {					// 遍历每个 segment，计算该向量在每个 segment 中的距离，加和成为 向量 y 与当前 database vector 的近似距离
				candidate += dists[comp_i * NCENT + pqcode[comp_i]];
		}
		bh->push(pqcode_i, candidate);
	}
	float bh_max = bh->max();
	// 遍历 partition 中的 keep 个向量，计算距离，出现更大的距离后才更新 bh_max
	for(unsigned pqcode_i = bh->capacity(); pqcode_i < n; ++pqcode_i) {
		float candidate = 0;
		const std::uint8_t* const pqcode = partition + pqcode_i * NSQ;
		for(int comp_i = 0; comp_i < NSQ; ++comp_i) {
				candidate += dists[comp_i * NCENT + pqcode[comp_i]];
		}
		if(candidate < bh_max) {
			bh->push(pqcode_i, candidate);
			bh_max = bh->max();
		}
	}
	return bh->max();													// 返回当前 binheap 中的最大值，也就是整个数据集中前 KEEP 个向量与 查询向量 y 的距离的最大值
}

2. Q127

Q127 是一个 SQ 压缩函数：

FORCE_INLINE
static inline std::int8_t Q127(const float x, const float min, const float max) {
	const float ratio = ((x-min)/(max-min));
	const float ratio_saturated = ratio < 1.0f ? ratio : 1.0f;
	return (static_cast<std::int8_t>(ratio_saturated*127));
}

可以看到它将 x 压缩到了 min-max 区间 128 段对应的整数值，也就是 (static_caststd::int8_t(ratio_saturated*127)) 。

3. min4_small_table

min4_small_table 是对前 4 个 segment 进行压缩的函数，我把值得注意的地方都写在了注释里：

FORCE_INLINE
static inline float min4_small_table(float qmax, const float* dists, __m128i (&min4)[4]) {
	float groups_min[4][16];
	float qmin = FLT_MAX;
	for(int sq_i = 0; sq_i < 4; ++sq_i) {							// 对前 4 个 segment 进行压缩， sq_i 指的是 segment 的索引
		// Init group minimums
		for(int group_i = 0; group_i < 16; ++group_i) {				// 每个 segment 中有 16 个局部距离组别， group_i 指的是局部距离分组的索引
			groups_min[sq_i][group_i] = FLT_MAX;					// 初始化每个 segment 中的 16 个局部距离的最小值为无穷大
		}
		// Find minimum for each group
		for(int cent_i = 0; cent_i < 256; ++cent_i) {
			const int group_i = cent_i % 16;						// 计算当前局部距离的分组索引，每个 segment 中有 256 个局部距离，分成 16 组，（0,16,32...256）分成一组，id为1
			const float candidate = dists[sq_i * NCENT + cent_i];
			if(candidate < groups_min[sq_i][group_i]) {				// 如果当前局部距离小于该组的最小值，则更新该组的最小值
				groups_min[sq_i][group_i] = candidate;
				if(candidate < qmin) {
					qmin = candidate;
				}
			}
		}
	}
	// 对后 4 个 segment 进行遍历，找到整个 LUTs 中最小的局部距离为 qmin ，后续压缩使用
	for(int sq_i = 4; sq_i < 8; ++sq_i) {							
		for(int cent_i = 0; cent_i < 256; ++cent_i) {
			const float candidate = dists[sq_i * NCENT + cent_i];
			if(candidate < qmin) {
				qmin = candidate;
			}
		}
	}
	// Quantize found minimums
	for(int sq_i = 0; sq_i < 4; ++sq_i) {
		min4[sq_i] = _mm_set_epi8(
				Q127(groups_min[sq_i][15], qmin, qmax),
				Q127(groups_min[sq_i][14], qmin, qmax),
				Q127(groups_min[sq_i][13], qmin, qmax),
				Q127(groups_min[sq_i][12], qmin, qmax),
				Q127(groups_min[sq_i][11], qmin, qmax),
				Q127(groups_min[sq_i][10], qmin, qmax),
				Q127(groups_min[sq_i][9], qmin, qmax),
				Q127(groups_min[sq_i][8], qmin, qmax),
				Q127(groups_min[sq_i][7], qmin, qmax),
				Q127(groups_min[sq_i][6], qmin, qmax),
				Q127(groups_min[sq_i][5], qmin, qmax),
				Q127(groups_min[sq_i][4], qmin, qmax),
				Q127(groups_min[sq_i][3], qmin, qmax),
				Q127(groups_min[sq_i][2], qmin, qmax),
				Q127(groups_min[sq_i][1], qmin, qmax),
				Q127(groups_min[sq_i][0], qmin, qmax)			    // 对 groups_min[sq_i][0] 进行 SQ 压缩，把 float 形的距离压缩成 int8 形的距离
				);

	}
	return qmin;
}

min4_small_table 的核心是，把每个分组中最小的局部距离找到，然后用这个局部距离代替整个分组中所有局部距离。在此基础上通过一个 SQ 压缩，把整个 segment 的全部 16 个分组，从原来的 16 * 32 = 512bits (16个最小值，最小值用 float32 表示)，压缩成了 16 * 8 = 128bits (最小值压缩成了 int8) ，这样每个 segment 的全部局部距离都可以被加载进一个 SIMD 寄存器，即使这个 SIMD 寄存器是最小的（不同型号的 SIMD 寄存器容量不一样，比如 AVX512 指令集就需要 SIMD 寄存器容量为 512 位）

4. ft4_small_table

ft4_small_table 是一个针对后 4 个 segments 的压缩算法，也就是针对我们的分组依据的 4 个 segment中的局部距离的压缩算法。因为分组操作，我们顺序扫描的向量的前 4 个 segment 中需要访问的短区间都属于同一个，所以我们就可以将这一部分的 16 个局部距离都压缩到一个 SIMD 寄存器中，避免从过远的内存中加载数据。

此外，之前的 min4_small_table 的压缩比非常极致，所以不可避免会出现较大的精度损失，一旦损失太大，所有的估计距离可能都会是下界，因此这个过滤算法就失效了。所以 ft4_small_table 通过使用更小的压缩比，保留更优秀的数据精度，来避免过滤失效。

FORCE_INLINE
static inline void ft4_small_table(float qmax, const float* dists, __m128i (&ft4)[4][16], const float qmin) {
	for(int sq_i = 0; sq_i < 4; ++sq_i) {
		const float* const sq_dists = dists + (sq_i + 4) * NCENT;		// 读取后 4 个 segment 的局部距离，跳过前 4 个 segment
		for(int h_cent_i = 0; h_cent_i < 16; ++h_cent_i) {
			const float* h_dists = sq_dists + h_cent_i * 16;			// 批量压缩，一次性 16 个 float32_t 类型的局部距离，
			ft4[sq_i][h_cent_i] = _mm_set_epi8(
					Q127(h_dists[15], qmin, qmax),
					Q127(h_dists[14], qmin, qmax),
					Q127(h_dists[13], qmin, qmax),
					Q127(h_dists[12], qmin, qmax),
					Q127(h_dists[11], qmin, qmax),
					Q127(h_dists[10], qmin, qmax),
					Q127(h_dists[9], qmin, qmax),
					Q127(h_dists[8], qmin, qmax),
					Q127(h_dists[7], qmin, qmax),
					Q127(h_dists[6], qmin, qmax),
					Q127(h_dists[5], qmin, qmax),
					Q127(h_dists[4], qmin, qmax),
					Q127(h_dists[3], qmin, qmax),
					Q127(h_dists[2], qmin, qmax),
					Q127(h_dists[1], qmin, qmax),
					Q127(h_dists[0], qmin, qmax)						// 依次压缩每个 float32 类型的局部距离，压缩成 int8_t 类型的局部距离
					);
		}
	}
}

可以看到，后面 4 个 segment ，每个 segment 原本有 256 个 float32 数据作为局部距离查询表。在对每个 float32 都进行 SQ 压缩后，变成了 256 个 int8 。然后每 16 个压缩局部距离就可以被放进同一个 SIMD 寄存器，也就是 ft4[sq_i][h_cent_i] = _mm_set_epi8() 函数做的事情。

需要额外注意的是，量化的上界和下界是不同的，上界是在整个 database vector 选择的前百分之 keep 的向量的 PQ Scan 得到的距离，而下界是局部距离的最小值。通过这种不对称的区间设计，让量化区间的上下界差距更大，从而将单个压缩距离的量化结果更集中于较小区间，保障了只要局部距离的加和大于 127 （量化编码的上界），就必定没用。见下图：

alt text

5. fast_scan_1

fast_scan_1 是查询入口函数。必须要注意的是，这里的距离计算扫描，都是默认输入的向量都属于同一个 group 。

fast_scan_1 后四个 segment 的查表距离加和：

         comp_block_size = 16;                       // 每次批量计算 16 个向量的距离
const __m128i low_bits_mask = _mm_set_epi64x(0x0f0f0f0f0f0f0f0f,
		0x0f0f0f0f0f0f0f0f);

         // 取出即将访问的每个 segment 的对应压缩数据到 SIMD 寄存器
         __m128i ft4_group[4];
         ft4_group[0] = ft4[0][hdr->values[0] >> 4];
         ft4_group[1] = ft4[1][hdr->values[1] >> 4];
         ft4_group[2] = ft4[2][hdr->values[2] >> 4];
         ft4_group[3] = ft4[3][hdr->values[3] >> 4];

// Components 4-5
__m128i comps_45 = _mm_loadu_si128(
		reinterpret_cast<const __m128i *>(partition
				+ 4 * comp_block_size)); 
const __m128i masked_comps_4 = _mm_and_si128(comps_45, low_bits_mask);      // 因为组别已经通过 group->value 确定了，代表索引的 8bits 的前面 4bits 都是没用的（这 16 个向量的这个segment都将访问同一个 group ），所以 _mm_and_si128 一下，让取指操作只用到后面 4 bits 作为偏移量
const __m128i partial_4 = _mm_shuffle_epi8(ft4_group[0], masked_comps_4);   // 取出每个向量的局部 int8 压缩编码
candidates = _mm_adds_epi8(candidates, partial_4);                          // 16 路 加和

comps_45 = _mm_srli_epi64(comps_45, 4);                                     // 因为只有 4 bits 索引是有用的，所以 8bit 可以同时存储两个 segment 的索引，向右移动4位，再用相同的掩码取低四位，便取出来了对应的 4bits 索引
const __m128i masked_comps_5 = _mm_and_si128(comps_45, low_bits_mask);
const __m128i partial_5 = _mm_shuffle_epi8(ft4_group[1], masked_comps_5);   
candidates = _mm_adds_epi8(candidates, partial_5);                          // 16 路 加和   

// Components 6-7
__m128i comps_67 = _mm_loadu_si128(
		reinterpret_cast<const __m128i *>(partition
				+ 5 * comp_block_size));
const __m128i masked_comps_6 = _mm_and_si128(comps_67, low_bits_mask);
const __m128i partial_6 = _mm_shuffle_epi8(ft4_group[2], masked_comps_6);
candidates = _mm_adds_epi8(candidates, partial_6);

const __m128i comps_7 = _mm_srli_epi64(comps_67, 4);
const __m128i masked_comp_7 = _mm_and_si128(comps_7, low_bits_mask);
const __m128i partial_7 = _mm_shuffle_epi8(ft4_group[3], masked_comp_7);
candidates = _mm_adds_epi8(candidates, partial_7);

fast_scan_1 前四个 segment 的查表距离加和：

         comp_block_size = 16;                       // 每次批量计算 16 个向量的距离
const __m128i low_bits_mask = _mm_set_epi64x(0x0f0f0f0f0f0f0f0f,
		0x0f0f0f0f0f0f0f0f);

// Component 0
const __m128i comps_0 = _mm_loadu_si128(
		reinterpret_cast<const __m128i *>(partition));
const __m128i masked_comps_0 = _mm_and_si128(comps_0, low_bits_mask);
__m128i candidates = _mm_shuffle_epi8(min4[0], masked_comps_0);         // 直接使用后 4 位作为偏移量，原因是这里的 索引值没有被压缩，还是 8bits ，这四个子段与 group 无关。所以只需要找到自己对应索引值属于哪个 group ，也就是 index % 16 ，按位运算就是取后4位。
// Components 1..3
for (int comp_i = 1; comp_i < 4; ++comp_i) {
	const __m128i comps = _mm_loadu_si128(
			reinterpret_cast<const __m128i *>(partition
					+ comp_i * comp_block_size));
	const __m128i masked_comps = _mm_and_si128(comps, low_bits_mask);
	const __m128i partial = _mm_shuffle_epi8(min4[comp_i], masked_comps);
	candidates = _mm_adds_epi8(candidates, partial);
}

这里其实就是对每个距离表的编码做运算，找到距离编码对应的 group ，也就是每个编码 % 16 代表的 group 。然后使用 _mm_shuffle_epi8 ，把每个距离编码对应的 group 的量化距离加载到 candidates 和 partial 中，执行加和。

这样，就获得了 16 个向量的 8 个 segments 的局部距离之和，这 16 个局部距离被放在同一个 SIMD 寄存器中，每个距离都用 int8 来表示。

其中需要额外说一嘴的是，使用 minimum 表来做查询的 4 个 segment 的向量索引还是使用的 uint8 来表示的，原因是后面的精细计算部分，需要这部分的准确距离结果，而且因为是对后面的非 group segment 进行的压缩，每个局部距离索引都可能位于 0- 255 .

3.2.2 为什么要读源码？

论文中对整个查询的流程描述的并不到位，差点给我看吐了。。。。。。

论文价值

1、论文是一个学习 PQ 、 CPU cache 、 SIMD 这些基础知识的优秀论文，值得我们仔细阅读理解。

2、论文中提到的高性能计算的优化逻辑是值得我们学习的，包括：

1、memory accessing times
2、instructions amount 
3、cpu cycles
4、L1 loads

3、论文的 SIMD 的代码实现值得一读