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

前言

单指令流多数据流(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位做翻转(flips10000)即可。

基于AVX2的向量化过滤

在StarRocks的底层,过滤器(Filter)是一个预分配空间的、无符号8位整形数的向量,用于表示WHEREHAVING子句的真值,每一位的取值为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

晚安。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 206,839评论 6 482
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 88,543评论 2 382
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 153,116评论 0 344
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 55,371评论 1 279
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 64,384评论 5 374
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,111评论 1 285
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,416评论 3 400
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 37,053评论 0 259
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 43,558评论 1 300
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 36,007评论 2 325
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,117评论 1 334
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,756评论 4 324
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,324评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,315评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,539评论 1 262
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,578评论 2 355
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,877评论 2 345

推荐阅读更多精彩内容