浅谈SIMD、向量化处理及其在StarRocks中的应用

LittleMagics2024-08-18 16:25

前言

单指令流多数据流(SIMD)及其衍生出来的向量化处理技术已经有了相当的历史，并且也是高性能数据库、计算引擎、多媒体库等组件的标配利器。笔者在两年多前曾经做过一次有关该主题的内部Geek分享，但可能是由于这个topic离实际研发场景比较远，当时听者寥寥。昨晚翻看硬盘中存的各种资料，翻到了相关内容，遂整理出来，顺便添加一些新东西。

SIMD

SIMD即"single instruction, multiple data"的缩写，是Flynn分类法对计算机的四大分类之一。它本质上是采用一个控制器来控制多个处理器，同时对一组数据中的每一条分别执行相同的操作，从而实现空间上的并行性的技术。

可见，"单指令流"指的是同时只能执行一种操作，"多数据流"则指的是在一组同构的数据(通常称为vector，即向量)上进行操作，如下图所示，其中PU = processing unit。

SIMD在现代计算机体系中的应用十分广泛，最典型的则是在GPU的像素处理流水线中。举个例子，如果要更改一整幅图像的亮度，只需要取出各像素的RGB值存入向量单元(向量单元很宽，可以存储多个像素的数据)，再同时将它们做相同的加减操作即可，效率很高。SIMD和MIMD流水线是GPU微架构的基础，就不再展开聊了。

那么CPU是如何实现SIMD的呢？答案是扩展指令集。Intel的第一版SIMD扩展指令集称为MMX，于1997年发布。后来至今的改进版本有SSE(Streaming SIMD Extensions)、AVX(Advanced Vector Extensions)，以及AMD的3DNow!等。我们可以通过cpuid类软件获得处理器对SIMD扩展指令集的支持信息，例如随便找一台服务器，执行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或者服务器处理器才支持SIMD扩展指令集，下图以笔者家用游戏PC中的AMD锐龙9 7950X3D处理器为例，可见同样支持。

下面简要介绍SSE指令集。

SSE指令集

SSE指令集是MMX的继任者，其第一版早在Pentium III时代就被引入了。随着新指令的扩充，又有了SSE2、SSE3、SSSE3、SSE4(包含4.1和4.2)等新版本。

SSE指令集以8个128位寄存器为基础，命名为XMM0~XMM7。在AMD64(即64位扩展)指令集中，又新增了XMM8~XMM15。一个XMM寄存器原本只能存储一种数据类型，即4个32位单精度浮点数，后来SSE2又扩展到能够存储以下类型：

  • 2个64位双精度浮点数
  • 2个64位 / 4个32位 / 8个16位整数
  • 16个字节或字符

SIMD指令分为两大类，一是标量(scalar)指令，二是打包(packed)指令。标量指令只对XMM寄存器中的最低位数据进行计算，打包指令则是对所有数据进行计算。下图示出SSE1中，单精度浮点数乘法的标量和打包运算。

观察指令助记符，mul表示乘法，接下来的s表示标量，p表示打包，最后一个s则表示类型为单精度浮点数(single-precision)。由图也可以发现，打包指令才是真正SIMD的，而标量指令是SISD的。

再举个小栗子，如果我们要实现两个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

注意数据移动指令movaps中的a表示对齐(align)。第一条指令的意思就是通过[v1]直接寻址得到向量的起点，并分别按照0、4、8、16字节的偏移量写入XMM0寄存器的低到高四个域。在数据本身已经按照16字节对齐的情况下，调用这种指令效率非常高。从寄存器写入内存也是同理的，如下图。

除了存取和数学运算指令外，SSE还提供了常用的比较、位移、位运算、类型转换、预取等指令。由此可见，SIMD对于那些严重依赖流程控制(flow control heavy)的任务，即有大量分支、跳转和条件判断的任务则不太适用。也就是说，SIMD主要被用来优化可并行计算的简单场景，以及可能被频繁调用的基础逻辑。

接下来再快速看一眼AVX指令集。

AVX指令集

AVX指令集是基于SSE指令集的扩展，在Sandy Bridge时代提出，Haswell时代又新增了AVX2。AVX指令集支持的数据类型与SSE本质上相同，但寄存器宽度翻了一倍，由128位来到了256位，称为YMM寄存器(SSE的XMM寄存器可以视作是YMM的低128位)，如下图所示。

以下是vhaddpd指令的图示，它分别将两个YMM寄存器中的64位浮点数水平相加(d代表double)，然后将结果交错存入第三个YMM寄存器中。

相比SSE，AVX支持更高效的位重排、三操作数指令(如上，即C = A + B)、非对齐访存等特性。StarRocks的SIMD优化主要就是基于AVX2做的，所以在部署文档的第一步，就是检查部署环境的CPU是否支持AVX2指令集。

说了这么多，最后以StarRocks为例简单看看SIMD扩展指令集在实际工程中的运用。

StarRocks向量化处理示例

如何运用SIMD指令集呢？主要有以下3种方法：

  • 直接编写内嵌汇编语句；
  • 利用厂商提供的扩展库函数。Intel将这类函数统称为Intrinsics，官方提供的速查手册见这里；
  • 开启编译器的优化(如GCC/G++的-msse2、-mavx2等)，编译器会自动将符合条件的情景(最简单的如数组相加、矩阵相乘)编译为SIMD指令。

向量化处理涉及到大量的case by case优化，在StarRocks BE源码中随处可见。我们可以查找形如#ifdef __SSE2__的宏定义，或者根据手册查找Intrinsic函数对应的头文件，如AVX2的头文件是<immintrin.h>，以此类推。

下面选取两段示例代码简单分析。

基于SSE2的向量化大小写转换

先上代码。

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));
}

根据手册简要介绍一下代码中涉及到的Intrinsic函数：

  • _mm_loadu_si128(mem_addr)：从内存地址mem_addr处加载128位的整形数据；
  • _mm_storeu_si128(mem_addr, a)：将128位的整形数据a存入内存地址mem_addr处；
  • _mm_set1_epi8(a)：将8位整形数据a广播到128位，即写入16个a；
  • _mm_cmpgt_epi8(a, b)：按8位比较a和b两个128位整形数，若a的对应8位比b的对应8位大，则填充对应位为全1，否则填充全0；
  • _mm_and_si128(a, b)和_mm_xor_si128(a, b)：两个128位数据之前按位与和按位异或。

由此可见，整个流程是一次性加载16个字符，然后并行判断字符是否在[A-Za-z]的范围内(注意掩码masks)，若符合条件，则根据大小写字母ASCII码值相差32的特性，对字符的第5位做翻转(flips为10000)即可。

基于AVX2的向量化过滤

在StarRocks的底层，过滤器(Filter)是一个预分配空间的、无符号8位整形数的向量，用于表示WHERE和HAVING子句的真值，每一位的取值为0或1，即表示为假或真。Filter和列(Column)是共生的，每种Column的实现都提供了对应的filter_range方法来过滤数据。以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;
}

还是根据手册简要介绍一下代码中涉及到的Intrinsic函数：

  • _mm256_setzero_si256()：返回一个256位的全0位图；
  • _mm256_loadu_si256(mem_addr)：从内存地址mem_addr处加载256位的整形数据；
  • _mm256_cmpgt_epi8(a, b)：按8位比较a和b两个256位整形数，若a的对应8位比b的对应8位大，则填充对应位为全1，否则填充全0；
  • _mm256_movemask_epi8(a)：根据256位整形数a的每个8位组的最高位生成掩码，一共32位长，返回一个int型结果。

由此可见，BE通过AVX2一次性加载一批32个真值进行判断。生成的掩码若为全0，表示全部不满足过滤条件，若为0xffffffff，则表示全部满足过滤条件，并拷贝结果。若是0、1混杂的情况，则调用内置的__builtin_ctz()函数取得掩码中末尾0的个数，然后直接跳过这些0位对应的数据，只将1对应的有效数据拷贝。如此循环，直到剩余的真值不满32个，循环遍历完即可。

The End

晚安。

热门推荐
01RAG 实践- Ollama+RagFlow 部署本地知识库02组基轨迹建模 GBTM的介绍与实现（Stata 或 R）032024年高教社杯数学建模国赛C题超详细解题思路分析04Coze扣子平台完整体验和实践（附国内和国际版对比）05苍穹外卖面试总结06【2024数模国赛赛题思路公开】国赛B题思路丨附可运行代码丨无偿自提07yolov8实战第五天——yolov8+ffmpg实时视频流检测并进行实时推流——（推流，保姆教学）08CCF-CSP认证考试 202406-3 文本分词 100分题解09【2024高教社杯全国大学生数学建模竞赛】B题 生产过程中的决策问题——解题思路 代码 论文10浏览器插件——ModHeader 的分享和学习