Press "Enter" to skip to content

27倍性能之旅 – 以大底库全库向量召回为例谈Profiling驱动的性能优化

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

问题

 

Problem

 

kNN(k Nearest Neighbor)定义
给定一个查询向量,按照某个选定的准则(如欧式距离),从底库中选择 个与查询向量最相似或者最邻近的向量的过程。

 

输入

 

查询向量(query):

 

底库(database): , 库中有 个向量,每个向量的维度为 ,即

 

输出

 

底库中与查询向量 点积 相似度最大的k个向量:

算法

 

Exhaustive Search, 穷搜

 

用例
只考虑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的性能发挥。

 

 

[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体系。

 


所以为了降低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;
}

 

例程:
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体系)转移到了读数据的端口。

 

 

Microcode Sequencer的bound效应开始显现。这个主要还是来自于我们计算方案中使用的hadd指令。下表给出了HADD指令的uOP数和port使用情况。

 

 

[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能不能有帮助。

 


下面试一下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驱动的优化的过程,我们在实际工作的问题可能比这个更复杂,但主要方法论还是通用的。另外,这里只说明了单核优化的情况,对多核优化其实并未涉及,但我个人认为只要在单核上把问题理解清楚了,多核扩展是相对比较自然的。

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注