27倍性能之旅 – 以大底庫全庫向量召回為例談Profiling驅動的性能優化

 

問題

  • Problem

    • kNN(k Nearest Neighbor)定義
      給定一個查詢向量，按照某個選定的準則（如歐式距離），從底庫中選擇個與查詢向量最相似或者最鄰近的向量的過程。

    • 輸入

      • 查詢向量(query):

      • 底庫(database): , 庫中有個向量，每個向量的維度為，即

    • 輸出

      • 底庫中與查詢向量點積相似度最大的k個向量：

         

    • 演算法

      • Exhaustive Search, 窮搜

  • 用例
    只考慮IP計算部分， 參數配置如下：

    參數 配置
    query batch size(bs) 32
    底庫大小(n) 1M條向量
    向量維度(d) 128
    向量數據精度(precision) FP32

    測試單核性能。

  • 程式碼 //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)分布。

Naive TMAM
看到這個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。

AVX2 Baseline TMAM

[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

[Impl #3] AVX-512 Query Batching

到現在為止，基於SIMD的計算優化就告一段落了。我們可以把目光投向Memory Bound。我們再仔細看一下Memory Bound的分布，可以看到幾乎全部都bound到了DRAM, L1/L2/L3 cache事實上失去了作用。

AVX3 Baseline Memory Bound Detail
出現這個狀況也是可解釋的，我們目前的實現方案可以如下圖表示。外層是query循環，內層是底庫循環，而且query的向量數是遠小於底庫的向量數的。我們可以從局部性的角度來打量query和資料庫向量在本計算方案中的情況：

  • query向量：有很好的空間局部性和時間局部性，query會連續與與底庫中每個向量計算，cache hierarchy 能夠很好的amortize掉memory read的overhead。

  • db向量：有很好的空間局部性，為SIMD的使用奠定基礎。但時間局部性非常差，每個向量被使用一次後，下一次使用就是在底庫被全庫遍歷後再回到該向量的時候了，而底庫又非常大，所以下一次使用的時候仍能命中cache的可能性為0。 因此，在cache behavior上來講變成了一個streaming數據，擊穿了cache體系。

Compute scheme
所以為了降低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。

AVX3 BP
這是一個非常有趣的VTune TMAM。首先看到如我們所料，我們通過增加db向量的時間局部性大幅降低了系統的Memory Bound(78.4% -> 5.8%)，基本達到我們的目的。進一步分析，可以發現

  • 系統的熱點轉移到了」Core Bound」, 各port都有比較大的壓力，特別是Port 2和Port 3這兩個壓力最大, 從下面的微架構圖中可以看出這兩個埠是Load/Store數據的埠，說明數據讀壓力還是較大，只不過這次壓力從數據源（記憶體/Cache體系）轉移到了讀數據的埠。

SKX Micro-architecture

  • Microcode Sequencer的bound效應開始顯現。這個主要還是來自於我們計算方案中使用的hadd指令。下表給出了HADD指令的uOP數和port使用情況。

HADD Instruction

[Impl #4] AVX-512 Block計算

沿著上一個section的分析，我們思考有沒有辦法降低core的port的壓力，並減少MicroSequencer的觸發率。
我們目前的計算方案，是每次直接計算一對vector的IP。以vector dimension 8為例，具體方案如下。從中我們可以看出，計算的冗餘度比較大從第二個hadd開始，有一半或以上的計算結果都不會倍最終結果用到，形成了一種「無效向量化」的效果，而且hadd指令的加入也使得MicroSequencer加入了指令的issue，拖慢了運行的效率。

Direct Scheme
Block優化可以消除計算冗餘度，同時消除hadd指令的使用。想法很簡單，把」1 v 1」的向量計算，變成」1 v N」的向量計算， 這裡N需要取一個AVX512暫存器里能保存的float的個數，也就是16。具體如下：

Block Scheme
按資料庫中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
下面試一下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的部分並沒有什麼影響。

Alt text

至此，優化基本告一段落。後面只有向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驅動的優化的過程，我們在實際工作的問題可能比這個更複雜，但主要方法論還是通用的。另外，這裡只說明了單核優化的情況，對多核優化其實並未涉及，但我個人認為只要在單核上把問題理解清楚了，多核擴展是相對比較自然的。

Tags:

VirMach 便宜 VPS

QNews

QNews