淺談SIMD、向量化處理及其在StarRocks中的應(yīng)用

前言

單指令流多數(shù)據(jù)流(SIMD)及其衍生出來(lái)的向量化處理技術(shù)已經(jīng)有了相當(dāng)?shù)臍v史,并且也是高性能數(shù)據(jù)庫(kù)、計(jì)算引擎、多媒體庫(kù)等組件的標(biāo)配利器。筆者在兩年多前曾經(jīng)做過(guò)一次有關(guān)該主題的內(nèi)部Geek分享,但可能是由于這個(gè)topic離實(shí)際研發(fā)場(chǎng)景比較遠(yuǎn),當(dāng)時(shí)聽(tīng)者寥寥。昨晚翻看硬盤(pán)中存的各種資料,翻到了相關(guān)內(nèi)容,遂整理出來(lái),順便添加一些新東西。

SIMD

SIMD即"single instruction, multiple data"的縮寫(xiě),是Flynn分類(lèi)法對(duì)計(jì)算機(jī)的四大分類(lèi)之一。它本質(zhì)上是采用一個(gè)控制器來(lái)控制多個(gè)處理器,同時(shí)對(duì)一組數(shù)據(jù)中的每一條分別執(zhí)行相同的操作,從而實(shí)現(xiàn)空間上的并行性的技術(shù)。

可見(jiàn),“單指令流”指的是同時(shí)只能執(zhí)行一種操作,“多數(shù)據(jù)流”則指的是在一組同構(gòu)的數(shù)據(jù)(通常稱(chēng)為vector,即向量)上進(jìn)行操作,如下圖所示,其中PU = processing unit。

SIMD在現(xiàn)代計(jì)算機(jī)體系中的應(yīng)用十分廣泛,最典型的則是在GPU的像素處理流水線中。舉個(gè)例子,如果要更改一整幅圖像的亮度,只需要取出各像素的RGB值存入向量單元(向量單元很寬,可以存儲(chǔ)多個(gè)像素的數(shù)據(jù)),再同時(shí)將它們做相同的加減操作即可,效率很高。SIMD和MIMD流水線是GPU微架構(gòu)的基礎(chǔ),就不再展開(kāi)聊了。

那么CPU是如何實(shí)現(xiàn)SIMD的呢?答案是擴(kuò)展指令集。Intel的第一版SIMD擴(kuò)展指令集稱(chēng)為MMX,于1997年發(fā)布。后來(lái)至今的改進(jìn)版本有SSE(Streaming SIMD Extensions)、AVX(Advanced Vector Extensions),以及AMD的3DNow!等。我們可以通過(guò)cpuid類(lèi)軟件獲得處理器對(duì)SIMD擴(kuò)展指令集的支持信息,例如隨便找一臺(tái)服務(wù)器,執(zhí)行cat /proc/cpuinfo命令,觀察flags域,如下。

processor   : 63
vendor_id   : GenuineIntel
cpu family  : 6
model       : 79
model name  : Intel(R) Xeon(R) CPU E5-2683 v4 @ 2.10GHz
stepping    : 1
microcode   : 0xb000040
cpu MHz     : 1272.637
cache size  : 40960 KB
physical id : 1
siblings    : 32
core id     : 15
cpu cores   : 16
apicid      : 63
initial apicid  : 63
fpu     : yes
fpu_exception   : yes
cpuid level : 20
wp      : yes
flags       : fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch epb cat_l3 cdp_l3 invpcid_single intel_pt ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a rdseed adx smap xsaveopt cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts md_clear spec_ctrl intel_stibp flush_l1d
bogomips    : 4204.62
clflush size    : 64
cache_alignment : 64
address sizes   : 46 bits physical, 48 bits virtual
power management:

并不僅有Intel或者服務(wù)器處理器才支持SIMD擴(kuò)展指令集,下圖以筆者家用游戲PC中的AMD銳龍9 7950X3D處理器為例,可見(jiàn)同樣支持。

下面簡(jiǎn)要介紹SSE指令集。

SSE指令集

SSE指令集是MMX的繼任者,其第一版早在Pentium III時(shí)代就被引入了。隨著新指令的擴(kuò)充,又有了SSE2、SSE3、SSSE3、SSE4(包含4.1和4.2)等新版本。

SSE指令集以8個(gè)128位寄存器為基礎(chǔ),命名為XMM0~XMM7。在AMD64(即64位擴(kuò)展)指令集中,又新增了XMM8~XMM15。一個(gè)XMM寄存器原本只能存儲(chǔ)一種數(shù)據(jù)類(lèi)型,即4個(gè)32位單精度浮點(diǎn)數(shù),后來(lái)SSE2又?jǐn)U展到能夠存儲(chǔ)以下類(lèi)型:

  • 2個(gè)64位雙精度浮點(diǎn)數(shù)
  • 2個(gè)64位 / 4個(gè)32位 / 8個(gè)16位整數(shù)
  • 16個(gè)字節(jié)或字符

SIMD指令分為兩大類(lèi),一是標(biāo)量(scalar)指令,二是打包(packed)指令。標(biāo)量指令只對(duì)XMM寄存器中的最低位數(shù)據(jù)進(jìn)行計(jì)算,打包指令則是對(duì)所有數(shù)據(jù)進(jìn)行計(jì)算。下圖示出SSE1中,單精度浮點(diǎn)數(shù)乘法的標(biāo)量和打包運(yùn)算。

觀察指令助記符,mul表示乘法,接下來(lái)的s表示標(biāo)量,p表示打包,最后一個(gè)s則表示類(lèi)型為單精度浮點(diǎn)數(shù)(single-precision)。由圖也可以發(fā)現(xiàn),打包指令才是真正SIMD的,而標(biāo)量指令是SISD的。

再舉個(gè)小栗子,如果我們要實(shí)現(xiàn)兩個(gè)4維向量v1和v2的加法,只需要三條SSE指令就夠了。

 movaps xmm0, [v1] ;xmm0 = v1.w | v1.z | v1.y | v1.x 
 addps xmm0, [v2]  ;xmm0 = v1.w+v2.w | v1.z+v2.z | v1.y+v2.y | v1.x+v2.x
 movaps [vec_res]  ;xmm0

注意數(shù)據(jù)移動(dòng)指令movaps中的a表示對(duì)齊(align)。第一條指令的意思就是通過(guò)[v1]直接尋址得到向量的起點(diǎn),并分別按照0、4、8、16字節(jié)的偏移量寫(xiě)入XMM0寄存器的低到高四個(gè)域。在數(shù)據(jù)本身已經(jīng)按照16字節(jié)對(duì)齊的情況下,調(diào)用這種指令效率非常高。從寄存器寫(xiě)入內(nèi)存也是同理的,如下圖。

除了存取和數(shù)學(xué)運(yùn)算指令外,SSE還提供了常用的比較、位移、位運(yùn)算、類(lèi)型轉(zhuǎn)換、預(yù)取等指令。由此可見(jiàn),SIMD對(duì)于那些嚴(yán)重依賴(lài)流程控制(flow control heavy)的任務(wù),即有大量分支、跳轉(zhuǎn)和條件判斷的任務(wù)則不太適用。也就是說(shuō),SIMD主要被用來(lái)優(yōu)化可并行計(jì)算的簡(jiǎn)單場(chǎng)景,以及可能被頻繁調(diào)用的基礎(chǔ)邏輯。

接下來(lái)再快速看一眼AVX指令集。

AVX指令集

AVX指令集是基于SSE指令集的擴(kuò)展,在Sandy Bridge時(shí)代提出,Haswell時(shí)代又新增了AVX2。AVX指令集支持的數(shù)據(jù)類(lèi)型與SSE本質(zhì)上相同,但寄存器寬度翻了一倍,由128位來(lái)到了256位,稱(chēng)為YMM寄存器(SSE的XMM寄存器可以視作是YMM的低128位),如下圖所示。

以下是vhaddpd指令的圖示,它分別將兩個(gè)YMM寄存器中的64位浮點(diǎn)數(shù)水平相加(d代表double),然后將結(jié)果交錯(cuò)存入第三個(gè)YMM寄存器中。

相比SSE,AVX支持更高效的位重排、三操作數(shù)指令(如上,即C = A + B)、非對(duì)齊訪存等特性。StarRocks的SIMD優(yōu)化主要就是基于AVX2做的,所以在部署文檔的第一步,就是檢查部署環(huán)境的CPU是否支持AVX2指令集。

說(shuō)了這么多,最后以StarRocks為例簡(jiǎn)單看看SIMD擴(kuò)展指令集在實(shí)際工程中的運(yùn)用。

StarRocks向量化處理示例

如何運(yùn)用SIMD指令集呢?主要有以下3種方法:

  • 直接編寫(xiě)內(nèi)嵌匯編語(yǔ)句;
  • 利用廠商提供的擴(kuò)展庫(kù)函數(shù)。Intel將這類(lèi)函數(shù)統(tǒng)稱(chēng)為Intrinsics,官方提供的速查手冊(cè)見(jiàn)這里;
  • 開(kāi)啟編譯器的優(yōu)化(如GCC/G++的-msse2、-mavx2等),編譯器會(huì)自動(dòng)將符合條件的情景(最簡(jiǎn)單的如數(shù)組相加、矩陣相乘)編譯為SIMD指令。

向量化處理涉及到大量的case by case優(yōu)化,在StarRocks BE源碼中隨處可見(jiàn)。我們可以查找形如#ifdef __SSE2__的宏定義,或者根據(jù)手冊(cè)查找Intrinsic函數(shù)對(duì)應(yīng)的頭文件,如AVX2的頭文件是<immintrin.h>,以此類(lèi)推。

下面選取兩段示例代碼簡(jiǎn)單分析。

基于SSE2的向量化大小寫(xiě)轉(zhuǎn)換

先上代碼。

template <char CA, char CZ>
static inline void vectorized_toggle_case(const Bytes* src, Bytes* dst) {
    const size_t size = src->size();
    // resize of raw::RawVectorPad16 is faster than std::vector because of
    // no initialization
    static_assert(sizeof(Bytes::value_type) == 1, "Underlying element type must be 8-bit width");
    static_assert(std::is_trivially_destructible_v<Bytes::value_type>,
                  "Underlying element type must have a trivial destructor");
    Bytes buffer;
    buffer.resize(size);
    uint8_t* dst_ptr = buffer.data();
    char* begin = (char*)(src->data());
    char* end = (char*)(begin + size);
    char* src_ptr = begin;
#if defined(__SSE2__)
    static constexpr int SSE2_BYTES = sizeof(__m128i);
    const char* sse2_end = begin + (size & ~(SSE2_BYTES - 1));
    const auto a_minus1 = _mm_set1_epi8(CA - 1);
    const auto z_plus1 = _mm_set1_epi8(CZ + 1);
    const auto flips = _mm_set1_epi8(32);

    for (; src_ptr > sse2_end; src_ptr += SSE2_BYTES, dst_ptr += SSE2_BYTES) {
        auto bytes = _mm_loadu_si128((const __m128i*)src_ptr);
        // the i-th byte of masks is set to 0xff if the corresponding byte is
        // between a..z when computing upper function (A..Z when computing lower function),
        // otherwise set to 0;
        auto masks = _mm_and_si128(_mm_cmpgt_epi8(bytes, a_minus1), _mm_cmpgt_epi8(z_plus1, bytes));
        // only flip 5th bit of lowcase(uppercase) byte, other bytes keep verbatim.
        _mm_storeu_si128((__m128i*)dst_ptr, _mm_xor_si128(bytes, _mm_and_si128(masks, flips)));
    }
#endif
    // only flip 5th bit of lowcase(uppercase) byte, other bytes keep verbatim.
    // i.e.  'a' and 'A' are 0b0110'0001 and 0b'0100'0001 respectively in binary form,
    // whether 'a' to 'A' or 'A' to 'a' conversion, just flip 5th bit(xor 32).
    for (; src_ptr < end; src_ptr += 1, dst_ptr += 1) {
        *dst_ptr = *src_ptr ^ (((CA <= *src_ptr) & (*src_ptr <= CZ)) << 5);
    }
    // move semantics
    dst->swap(reinterpret_cast<Bytes&>(buffer));
}

根據(jù)手冊(cè)簡(jiǎn)要介紹一下代碼中涉及到的Intrinsic函數(shù):

  • _mm_loadu_si128(mem_addr):從內(nèi)存地址mem_addr處加載128位的整形數(shù)據(jù);
  • _mm_storeu_si128(mem_addr, a):將128位的整形數(shù)據(jù)a存入內(nèi)存地址mem_addr處;
  • _mm_set1_epi8(a):將8位整形數(shù)據(jù)a廣播到128位,即寫(xiě)入16個(gè)a
  • _mm_cmpgt_epi8(a, b):按8位比較a和b兩個(gè)128位整形數(shù),若a的對(duì)應(yīng)8位比b的對(duì)應(yīng)8位大,則填充對(duì)應(yīng)位為全1,否則填充全0;
  • _mm_and_si128(a, b)_mm_xor_si128(a, b):兩個(gè)128位數(shù)據(jù)之前按位與和按位異或。

由此可見(jiàn),整個(gè)流程是一次性加載16個(gè)字符,然后并行判斷字符是否在[A-Za-z]的范圍內(nèi)(注意掩碼masks),若符合條件,則根據(jù)大小寫(xiě)字母ASCII碼值相差32的特性,對(duì)字符的第5位做翻轉(zhuǎn)(flips10000)即可。

基于AVX2的向量化過(guò)濾

在StarRocks的底層,過(guò)濾器(Filter)是一個(gè)預(yù)分配空間的、無(wú)符號(hào)8位整形數(shù)的向量,用于表示WHEREHAVING子句的真值,每一位的取值為0或1,即表示為假或真。Filter和列(Column)是共生的,每種Column的實(shí)現(xiàn)都提供了對(duì)應(yīng)的filter_range方法來(lái)過(guò)濾數(shù)據(jù)。以BinaryColumnBase為例,其filter_range方法的源碼如下。

template <typename T>
size_t BinaryColumnBase<T>::filter_range(const Filter& filter, size_t from, size_t to) {
    auto start_offset = from;
    auto result_offset = from;

    uint8_t* data = _bytes.data();

#ifdef __AVX2__
    const uint8_t* f_data = filter.data();

    int simd_bits = 256;
    int batch_nums = simd_bits / (8 * (int)sizeof(uint8_t));
    __m256i all0 = _mm256_setzero_si256();

    while (start_offset + batch_nums < to) {
        __m256i f = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(f_data + start_offset));
        uint32_t mask = _mm256_movemask_epi8(_mm256_cmpgt_epi8(f, all0));

        if (mask == 0) {
            // all no hit, pass
        } else if (mask == 0xffffffff) {
            // all hit, copy all

            // copy data
            T size = _offsets[start_offset + batch_nums] - _offsets[start_offset];
            memmove(data + _offsets[result_offset], data + _offsets[start_offset], size);

            // set offsets, try vectorized
            T* offset_data = _offsets.data();
            for (int i = 0; i < batch_nums; ++i) {
                // TODO: performance, all sub one same offset ?
                offset_data[result_offset + i + 1] = offset_data[result_offset + i] +
                                                     offset_data[start_offset + i + 1] - offset_data[start_offset + i];
            }

            result_offset += batch_nums;
        } else {
            // skip not hit row, it's will reduce compare when filter layout is sparse,
            // like "00010001...", but is ineffective when the filter layout is dense.

            uint32_t zero_count = Bits::CountTrailingZerosNonZero32(mask);
            uint32_t i = zero_count;
            while (i < batch_nums) {
                mask = zero_count < 31 ? mask >> (zero_count + 1) : 0;

                T size = _offsets[start_offset + i + 1] - _offsets[start_offset + i];
                // copy date
                memmove(data + _offsets[result_offset], data + _offsets[start_offset + i], size);

                // set offsets
                _offsets[result_offset + 1] = _offsets[result_offset] + size;
                zero_count = Bits::CountTrailingZeros32(mask);
                result_offset += 1;
                i += (zero_count + 1);
            }
        }
        start_offset += batch_nums;
    }
#endif

    for (auto i = start_offset; i < to; ++i) {
        if (filter[i]) {
            DCHECK_GE(_offsets[i + 1], _offsets[i]);
            T size = _offsets[i + 1] - _offsets[i];
            // copy data
            memmove(data + _offsets[result_offset], data + _offsets[i], size);

            // set offsets
            _offsets[result_offset + 1] = _offsets[result_offset] + size;

            result_offset++;
        }
    }

    this->resize(result_offset);
    return result_offset;
}

還是根據(jù)手冊(cè)簡(jiǎn)要介紹一下代碼中涉及到的Intrinsic函數(shù):

  • _mm256_setzero_si256():返回一個(gè)256位的全0位圖;
  • _mm256_loadu_si256(mem_addr):從內(nèi)存地址mem_addr處加載256位的整形數(shù)據(jù);
  • _mm256_cmpgt_epi8(a, b):按8位比較a和b兩個(gè)256位整形數(shù),若a的對(duì)應(yīng)8位比b的對(duì)應(yīng)8位大,則填充對(duì)應(yīng)位為全1,否則填充全0;
  • _mm256_movemask_epi8(a):根據(jù)256位整形數(shù)a的每個(gè)8位組的最高位生成掩碼,一共32位長(zhǎng),返回一個(gè)int型結(jié)果。

由此可見(jiàn),BE通過(guò)AVX2一次性加載一批32個(gè)真值進(jìn)行判斷。生成的掩碼若為全0,表示全部不滿(mǎn)足過(guò)濾條件,若為0xffffffff,則表示全部滿(mǎn)足過(guò)濾條件,并拷貝結(jié)果。若是0、1混雜的情況,則調(diào)用內(nèi)置的__builtin_ctz()函數(shù)取得掩碼中末尾0的個(gè)數(shù),然后直接跳過(guò)這些0位對(duì)應(yīng)的數(shù)據(jù),只將1對(duì)應(yīng)的有效數(shù)據(jù)拷貝。如此循環(huán),直到剩余的真值不滿(mǎn)32個(gè),循環(huán)遍歷完即可。

The End

晚安。

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書(shū)系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容