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 = ; i < bs; i++) {
   for (size_t j = ; 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 = ;
    for (size_t i = ; 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 = ; i < bs; i++) {
    for (size_t j = ; 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 = ; i < size; i += ) {
        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 + ),
                                   _mm256_loadu_ps(rhs + i + ),
                                   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, );
    __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 = ; i < bs; i++) {
    for (size_t j = ; 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 = ; i < dim; i += ) {
        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 + ),
                                   _mm512_loadu_ps(rhs + i + ),
                                   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), ));

    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总共需要 代码如下：

例程：
for (size_t i = ; 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(,  * bs * sizeof(__m512));
        zmm_sum1 = zmm_sum0 + bs;
    }
    memset(zmm_sum0, ,  * bs * sizeof(__m512));

    __m512 v_base0, v_base1;
    __m512 v_q0, v_q1;

    for (size_t i = ; i < dim; i += ) {
        v_base0 = _mm512_loadu_ps(base + i);
        v_base1 = _mm512_loadu_ps(base + i + );

        for (size_t j = ; j < bs; j++) {
            v_q0 = _mm512_loadu_ps(query + j * dim + i);
            v_q1 = _mm512_loadu_ps(query + j * dim + i + );

            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 = ; 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 = ; 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(,  * bs * sizeof(__m512));
        zmm_sum1 = zmm_sum0 + bs;
    }
    memset(zmm_sum0, ,  * bs * sizeof(__m512));

    __m512 v_base0, v_base1;
    __m512 v_q0, v_q1;

    for (size_t i = ; i < dim; i += ) {
        v_base0 = _mm512_loadu_ps(base + i);
        v_base1 = _mm512_loadu_ps(base + i + );

        for (size_t j = ; j < bs; j++) {
            v_q0 = _mm512_loadu_ps(query + j * dim + i);
            v_q1 = _mm512_loadu_ps(query + j * dim + i + );

            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 = ; 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条向量来算：从指令数上来讲，直接法需要 下面给出实现：

例程：
for (size_t i = ; i < base_vec_num; i += ) {
    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(, bs * sizeof(__m512));
    }
    memset(zmm_sum0, ,  * bs * sizeof(__m512));

    __m512 v_base0;
    __m512 v_q0, v_q1;

    for (size_t i = ; i < dim; i += ) {
        v_base0 = _mm512_loadu_ps(base + i * );

        for (size_t j = ; j < bs; j += ) {
            v_q0 = _mm512_extload_ps(query + j * dim + i);
            v_q1 = _mm512_extload_ps(query + (j + ) * dim + i);

            zmm_sum0[j] = _mm512_fmadd_ps(v_base0, v_q0, zmm_sum0[j]);
            zmm_sum0[j + ] = _mm512_fmadd_ps(v_base0, v_q1, zmm_sum0[j + ]);
        }
        _mm_prefetch(base + (i + ) * , _MM_HINT_T0);
    }

    for (size_t j = ; j < bs; j++) {
        _mm512_store_ps(scores + j * , 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(, bs * sizeof(__m512));
    }
    memset(zmm_sum0, ,  * bs * sizeof(__m512));

    __m512 v_base0;
    __m512 v_q0, v_q1;

    for (size_t i = ; i < dim; i += ) {
        v_base0 = _mm512_loadu_ps(base + i * );

        for (size_t j = ; j < bs; j += ) {
            v_q0 = _mm512_extload_ps(query + j * dim + i);
            v_q1 = _mm512_extload_ps(query + (j + ) * dim + i);

            zmm_sum0[j] = _mm512_fmadd_ps(v_base0, v_q0, zmm_sum0[j]);
            zmm_sum0[j + ] = _mm512_fmadd_ps(v_base0, v_q1,zmm_sum0[j + ]);
        }
        _mm_prefetch(base + (i + ) * , _MM_HINT_T0);
    }

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