@(Engineering Practice)[Fun For Deep Dive]
问题
-
Problem
- kNN(k Nearest Neighbor)定义
给定一个查询向量，按照某个选定的准则（如欧式距离），从底库中选择个与查询向量最相似或者最邻近的向量的过程。 - 输入
- 查询向量(query):
- 底库(database): , 库中有个向量，每个向量的维度为，即
- 输出
- 底库中与查询向量点积相似度最大的k个向量：
- 底库中与查询向量点积相似度最大的k个向量：
- 算法
- Exhaustive Search, 穷搜
- kNN(k Nearest Neighbor)定义
-
用例
只考虑IP计算部分， 参数配置如下：参数 配置 query batch size(bs) 32 底库大小(n) 1M条向量 向量维度(d) 128 向量数据精度(precision) FP32
测试单核性能。
- 代码
https://github.com/yao-matrix/mProto/tree/master/snippet/c/dp
实现
[Impl #0] Naive
最朴素的实现就是循环计算每个query向量与底库中每个向量的内积。实现如下：
例程：
for (size_t i = 0; i < bs; i++) {
for (size_t j = 0; j < base_vec_num; j++) {
dotp_naive(query + i * vec_dim, base + j * vec_dim,
vec_dim, &scores[i * base_vec_num + j]);
}
}
子例程：
static void dotp_naive(const float *lhs, const float *rhs,
size_t dim, float *score) {
float res = 0;
for (size_t i = 0; i < dim; i++) {
res += lhs[i] * rhs[i];
}
*score = res;
return;
}
性能显然是不好的，在至强CascadeLake 8269C上单核平均每batch需要5426.046 ms。
祭出 VTune大法后，得到如下的TMAM (Top-down Micro-architecture Analysis Method)分布。
看到这个TMAM图的时候，一定要先按捺住自己想去拯救Memory Bound那一片红海的冲动。优化的过程中，要牢牢守住“先计算后数据”的初心，因为数据往往是可以配合计算的，所以首先是要把计算这跟柱子立牢了，然后在这个基础上做数据优化。
计算优化的第一步是看SIMD有没有被用上：
- 我们看到整个workload retire的操作中，有38.4%的微操作是FP Arithmetic，即浮点算术操作，而这些操作目前全部都是标量实现的，没有SIMD化。
- 同时，通过Vector Capacity Usage(FPU) 的利用率 6.2%也可以看出CPU的向量计算能力没有被充分利用。
所以，虽然现在是Memory Bound，但我们还是不能“头疼医头，脚疼医脚”，而是应该先做向量化的优化，再来分析。
[Impl #1] AVX2 Vectorization
使用AVX-2来向量化IP运算，其余保持不变。代码如下：
例程：
for (size_t i = 0; i < bs; i++) {
for (size_t j = 0; j < base_vec_num; j++) {
dotp_avx2(query + i * vec_dim, base + j * vec_dim,
vec_dim, &scores[i * base_vec_num + j]);
}
}
子例程：
void dotp_avx2(const float *lhs, const float *rhs, size_t size, float *score) {
__m256 ymm_sum0 = _mm256_setzero_ps();
__m256 ymm_sum1 = _mm256_setzero_ps();
for (size_t i = 0; i < size; i += 16) {
ymm_sum0 = _mm256_fmadd_ps(_mm256_loadu_ps(lhs + i),
_mm256_loadu_ps(rhs + i),
ymm_sum0);
ymm_sum1 = _mm256_fmadd_ps(_mm256_loadu_ps(lhs + i + 8),
_mm256_loadu_ps(rhs + i + 8),
ymm_sum1);
}
*score = horizontal_add_v256(_mm256_add_ps(ymm_sum0, ymm_sum1));
return;
}
static inline float horizontal_add_v256(__m256 v) {
__m256 x1 = _mm256_hadd_ps(v, v);
__m256 x2 = _mm256_hadd_ps(x1, x1);
__m128 x3 = _mm256_extractf128_ps(x2, 1);
__m128 x4 = _mm_add_ss(_mm256_castps256_ps128(x2), x3);
return _mm_cvtss_f32(x4);
}
AVX2向量化带来了性能提升，在至强CascadeLake 8269C上单核平均每batch延时下降到了2143.29 ms。
同样，我们把VTune的TMAM数据拉出来端详端详。可以看到这下FP Arithmetic 都落到FP Vector上了，同时Vector Capacity Usage也提高到48.8%。只是Memory Bound更严重了，飙到了76.1%。这时要做的，还是按捺住自己想要去优化数据部分的小心脏，继续看看还有没有计算优化的空间。一个自然的想法是：试试AVX-512。
[Impl #2] AVX-512 Vectorization
使用AVX-512来向量化IP运算，其余保持不变。代码如下：
例程：
for (size_t i = 0; i < bs; i++) {
for (size_t j = 0; j < base_vec_num; j++) {
dotp_avx3(query + i * vec_dim, base + j * vec_dim,
vec_dim, &scores[i * base_vec_num + j]);
}
}
子例程：
void dotp_avx3(const float *lhs, const float *rhs, size_t dim, float *score) {
__m512 zmm_sum0 = _mm512_setzero_ps();
__m512 zmm_sum1 = _mm512_setzero_ps();
for (size_t i = 0; i < dim; i += 32) {
zmm_sum0 = _mm512_fmadd_ps(_mm512_loadu_ps(lhs + i),
_mm512_loadu_ps(rhs + i), zmm_sum0);
zmm_sum1 = _mm512_fmadd_ps(_mm512_loadu_ps(lhs + i + 16),
_mm512_loadu_ps(rhs + i + 16), zmm_sum1);
}
*score = horizontal_add_v512(_mm512_add_ps(zmm_sum0, zmm_sum1));
return;
}
static inline float horizontal_add_v512(__m512 v) {
__m256 low = _mm512_castps512_ps256(v);
__m256 high = _mm256_castpd_ps(_mm512_extractf64x4_pd(
_mm512_castps_pd(v), 1));
return horizontal_add_v256(low + high);
}
跑下来， 单核平均每batch延时进一步下降到了1589.99 ms。上TMAM，如下图。我们可以看到：
- 因为AVX-512的指令位宽是AVX2的加倍，所以处理相同的problem size所需要的指令数大约是AVX2的一半。这个可以通过比较Instructions Retired得到验证。
- 一半的指令数的减少并没有等比反映到延时性能中，我们可以看到最终延时仅有25.8%的降低。这里面有两个原因：
- AVX-512造成了CPU的降频，从Average CPU Frequency可以看到在AVX-512实现中频率从AVX2时的2.9 GHz降到了2.5 GHz, 降了0.4 GHz.
-
Memory Bound继续增大到了78.4%， 数据瓶颈阻碍了AVX-512的性能发挥。
AVX3 Baseline | center
[Impl #3] AVX-512 Query Batching
到现在为止，基于SIMD的计算优化就告一段落了。我们可以把目光投向Memory Bound。我们再仔细看一下Memory Bound的分布，可以看到几乎全部都bound到了DRAM, L1/L2/L3 cache事实上失去了作用。
出现这个状况也是可解释的，我们目前的实现方案可以如下图表示。外层是query循环，内层是底库循环，而且query的向量数是远小于底库的向量数的。我们可以从局部性的角度来打量query和数据库向量在本计算方案中的情况：
- query向量：有很好的空间局部性和时间局部性，query会连续与与底库中每个向量计算，cache hierarchy 能够很好的amortize掉memory read的overhead。
- db向量：有很好的空间局部性，为SIMD的使用奠定基础。但时间局部性非常差，每个向量被使用一次后，下一次使用就是在底库被全库遍历后再回到该向量的时候了，而底库又非常大，所以下一次使用的时候仍能命中cache的可能性为0。 因此，在cache behavior上来讲变成了一个streaming数据，击穿了cache体系。
Compute Scheme|center
所以为了降低Memory Bound，主要需要考虑如何提高db向量的时间局部性。一个自然的想法就是交换内外层循环。这样，每个db向量load进来后都会被使用多次，而不是像以前那样只有一次，这些计算amortize了一次db向量的memory read overhead，减少了DRAM的access。这里只有一个问题，交换循环后query的时间局部性会不会被牺牲掉？在这个问题中，我们算一下query总共需要 数据，而一个db向量需要, 两个加起来小于L1 DCache的size 32KB，所以query的cache局部性不会被牺牲掉。
代码如下：
例程：
for (size_t i = 0; i < base_vec_num; i++) {
dotp_avx3_qb(query, base + i * vec_dim, bs, vec_dim, scores + i * bs);
}
子例程（memalign的部分在正式实现时需要移到初始化API中，并在destroy代码中free内存。这里为简化代码忽略内存泄漏，不影响性能结论）：
void dotp_avx3_qb(const float *query, const float *base, size_t bs, size_t dim, float *scores) {
static __m512 *zmm_sum0 = NULL;
static __m512 *zmm_sum1 = NULL;
if (zmm_sum0 == NULL) {
zmm_sum0 = (__m512*)memalign(64, 2 * bs * sizeof(__m512));
zmm_sum1 = zmm_sum0 + bs;
}
memset(zmm_sum0, 0, 2 * bs * sizeof(__m512));
__m512 v_base0, v_base1;
__m512 v_q0, v_q1;
for (size_t i = 0; i < dim; i += 32) {
v_base0 = _mm512_loadu_ps(base + i);
v_base1 = _mm512_loadu_ps(base + i + 16);
for (size_t j = 0; j < bs; j++) {
v_q0 = _mm512_loadu_ps(query + j * dim + i);
v_q1 = _mm512_loadu_ps(query + j * dim + i + 16);
zmm_sum0[j] = _mm512_fmadd_ps(v_base0, v_q0, zmm_sum0[j]);
zmm_sum1[j] = _mm512_fmadd_ps(v_base1, v_q1, zmm_sum1[j]);
}
}
for (size_t j = 0; j < bs; j++) {
scores[j] = horizontal_add_v512(
_mm512_add_ps(zmm_sum0[j], zmm_sum1[j]));
}
// free(zmm_sum0);
// zmm_sum0 = NULL;
return;
}
单核平均每batch延时进一步从1589.99 ms下降到了348.05 ms， 性能提升了约3.5倍。这个时候再看一下VTune。
这是一个非常有趣的VTune TMAM。首先看到如我们所料，我们通过增加db向量的时间局部性大幅降低了系统的Memory Bound(78.4% -> 5.8%)，基本达到我们的目的。进一步分析，可以发现
-
系统的热点转移到了"Core Bound", 各port都有比较大的压力，特别是Port 2和Port 3这两个压力最大, 从下面的微架构图中可以看出这两个端口是Load/Store数据的端口，说明数据读压力还是较大，只不过这次压力从数据源（内存/Cache体系）转移到了读数据的端口。
SKX Micro-architecture|center -
Microcode Sequencer的bound效应开始显现。这个主要还是来自于我们计算方案中使用的hadd指令。下表给出了HADD指令的uOP数和port使用情况。
HADD Instruction | center
[Impl #4] AVX-512 Block计算
沿着上一个section的分析，我们思考有没有办法降低core的port的压力，并减少MicroSequencer的触发率。
我们目前的计算方案，是每次直接计算一对vector的IP。以vector dimension 8为例，具体方案如下。从中我们可以看出，计算的冗余度比较大从第二个hadd开始，有一半或以上的计算结果都不会倍最终结果用到，形成了一种“无效向量化”的效果，而且hadd指令的加入也使得MicroSequencer加入了指令的issue，拖慢了运行的效率。
Block优化可以消除计算冗余度，同时消除hadd指令的使用。想法很简单，把"1 v 1"的向量计算，变成"1 v N"的向量计算， 这里N需要取一个AVX512寄存器里能保存的float的个数，也就是16。具体如下：
按数据库中16条向量来算：从指令数上来讲，直接法需要条指令，而block优化需要条指令，指令数大大降低了，而且也消除了hadd。但是Block优化需要增加一个额外的操作，就是每16行db向量需要做一个转置，这个可以在数据库初始化的时候一次性做掉就可以了，不需要每次都做，所以这个时间是可以不算的。
下面给出实现：
例程：
for (size_t i = 0; i < base_vec_num; i += 16) {
dotp_avx3_block(query, base + i * vec_dim, bs, vec_dim, scores + i * bs);
}
子例程：
static inline __m512 _mm512_extload_ps(float const* daddr) {
__m128 d = _mm_load_ps(daddr);
return _mm512_broadcastss_ps(d);
}
void dotp_avx3_block(const float *query, const float *base,
size_t bs, size_t dim, float *scores) {
static __m512 *zmm_sum0 = NULL;
if (zmm_sum0 == NULL) {
zmm_sum0 = (__m512*)memalign(64, bs * sizeof(__m512));
}
memset(zmm_sum0, 0, 2 * bs * sizeof(__m512));
__m512 v_base0;
__m512 v_q0, v_q1;
for (size_t i = 0; i < dim; i += 1) {
v_base0 = _mm512_loadu_ps(base + i * 16);
for (size_t j = 0; j < bs; j += 2) {
v_q0 = _mm512_extload_ps(query + j * dim + i);
v_q1 = _mm512_extload_ps(query + (j + 1) * dim + i);
zmm_sum0[j] = _mm512_fmadd_ps(v_base0,
v_q0, zmm_sum0[j]);
zmm_sum0[j + 1] = _mm512_fmadd_ps(v_base0,
v_q1, zmm_sum0[j + 1]);
}
_mm_prefetch(base + (i + 1) * 16, _MM_HINT_T0);
}
for (size_t j = 0; j < bs; j++) {
_mm512_store_ps(scores + j * 16, zmm_sum0[j]);
}
// free(zmm_sum0);
// zmm_sum0 = NULL;
return;
}
跑下来，单核平均每batch延时进一步从348.05 ms下降到了206.53 ms， 性能提升了约1.7倍。再看一下VTune。可以看到port 2、port 3的压力下降了，Micro Sequencer的热点也不见了，FPU的利用率也达到了100%。另外观察到两个现象：
- 数据access端口压力还是最大，包括Port 2, Port 3, Port 4。尤其是Port 4是store端口，与Back-End Bound里面的Store Bound是相互映证的。目前还没有想到好的办法缓解。
-
L2 Bound比较significant, 另有3.7%的DRAM bound, 可以再试试prefetch能不能有帮助。
AVX3 Block|center
下面试一下prefetch, code如下。
子例程：
void dotp_avx3_block(const float *query, const float *base,
size_t bs, size_t dim, float *scores) {
static __m512 *zmm_sum0 = NULL;
if (zmm_sum0 == NULL) {
zmm_sum0 = (__m512*)memalign(64, bs * sizeof(__m512));
}
memset(zmm_sum0, 0, 2 * bs * sizeof(__m512));
__m512 v_base0;
__m512 v_q0, v_q1;
for (size_t i = 0; i < dim; i += 1) {
v_base0 = _mm512_loadu_ps(base + i * 16);
for (size_t j = 0; j < bs; j += 2) {
v_q0 = _mm512_extload_ps(query + j * dim + i);
v_q1 = _mm512_extload_ps(query + (j + 1) * dim + i);
zmm_sum0[j] = _mm512_fmadd_ps(v_base0,
v_q0, zmm_sum0[j]);
zmm_sum0[j + 1] = _mm512_fmadd_ps(v_base0,
v_q1, zmm_sum0[j + 1]);
}
_mm_prefetch(base + (i + 1) * 16, _MM_HINT_T0);
}
for (size_t j = 0; j < bs; j++) {
_mm512_store_ps(scores + j * 16, zmm_sum0[j]);
}
// free(zmm_sum0);
// zmm_sum0 = NULL;
return;
}
跑下来，单核平均每batch延时进一步从206.53 ms下降到了195.78 ms，约5%的性能提升。这个prefetch是prefetch到L1的，从TMAM分析可以看到，prefetch把DRAM bound消除了，但对L2 Bound的部分并没有什么影响。
至此，优化基本告一段落。后面只有向hardware提需求了， 如增大L1 cache size, 增加store port等。
总结与思考
我们回顾这次优化之旅，效果还是比较显著的：
| 方案 | 单核batch延时（单位：ms） |
|---|---|
| Naive | 5426.046 |
| AVX2 Baseline | 2143.29 |
| AVX512 Baselie | 1589.99 |
| AVX512 Query Batching | 348.05 |
| AVX512 Block | 206.53 |
| AVX512 Block + Prefetch | 195.78 |
一路下来，我们把单核batch延时从5426.046 ms下降到了195.78 ms, 性能提升了27倍。
在优化过程中，我们也收获不少思考：
- 编译器自动向量化效率比较低。在Naive算法中，我们在编译选项中已经开了native优化的选项，但从结果中我们可以看到，其向量化的效率并不理想，目前还是需要人工向量化。
- Profiling驱动优化是比较有效的优化方法。通过紧密结合profiling数据和对workload的理解，我们可以一步一步打磨workload和计算平台的适切性，从而获得较好的性能增益。
- 优化过程中要坚持“计算第一，数据配合”的原则，以计算优化为首要优先级，待计算优化稳定后再配合数据优化。
在本文中，我们对一个简单的问题进行了尝试，这是为了比较容易去说明profiling驱动的优化的过程，我们在实际工作的问题可能比这个更复杂，但主要方法论还是通用的。另外，这里只说明了单核优化的情况，对多核优化其实并未涉及，但我个人认为只要在单核上把问题理解清楚了，多核扩展是相对比较自然的。