论整数转 bool array 的高效实现

IceBear98
719 阅读8分钟

问题：如何高效地实现整数转成 bool array?

函数原型：

std::array<bool, 8> bitarray(unsigned char);

使用效果：

for (auto b : bitarray(0b11100100) {
    printf("%d ", b);
}

// 输出（高位在前，低位在后）
1 1 1 0 0 1 0 0

相信很多人对这个问题深不以为然，这不就是入门级的位运算练习题嘛。这个都不用分分钟了，秒秒钟就好的：

#include <stdio.h>
#include <array>

std::array<bool, 8>
bitarray(unsigned char x)
{
    std::array<bool, 8> r;
    for (int i = 0; i < 8; ++i) {
        r[i] = static_cast<bool>((x >> (7 - i)) & 1);
    }
    return r;
}

int main()
{
    for (auto b : bitarray(0b11100100)) {
        printf("%d ", b);
    }
}

可是，这种写法真的高效吗？

甚是不幸的是，如果你查看一下这种写法生成的汇编，你将会感到非常失望。而且更加沮丧的是，C++四大编译器，gcc clang msvc icc，居然都没法有效对如此简单常见的代码有任何行之有效的优化。

（当然，这些都是 2021 年年中的编译器的水平，如果在未来编译优化做的更智能了，你也不要做时空警察跨时空执法）

这是 g++12，开 -O2 -march=native 的效果：

g++ -O2 -march=native

image.png

这是开到了 -O3 的效果，8 次循环被展开：

g++ -O3 -march=native

image.png

这是 clang13 开 -O2 -march=native 的效果，可见其中用到了 sse 系列的指令。

clang -O2 -march=native

image.png

这还只是部分，太长了，没法截全。

这是 msvc2019 在 O2 下的效果：

image.png

这是 intel 自家的 icc2021.2.0 编译器在 -O2 -march=native 下的效果，可见其中用到了更为高级的 avx 系列指令：

image.png

而且，值得大家注意的是，笔者在以上若干例中都开了 -march=native 选项。笔者工作的计算机是 skylake 架构，这意味绝大多数主流的指令集，除了最新的 avx512 外，都已经开启，允许编译器使用它们去优化代码了。但，效果真的不尽如人意。

有什么优化方案吗？

相信稍有优化经验的猿们都会想到使用 SIMD，在指令层面将循环并行化。

即原来写法中的循环：

    for (int i = 0; i < 8; ++i) {
        r[i] = static_cast<bool>((x >> (7 - i)) & 1);
    }

展开为：

    r[0] = static_cast<bool>((x >> 7) & 1);
    r[1] = static_cast<bool>((x >> 6) & 1);
    r[2] = static_cast<bool>((x >> 5) & 1);
    r[3] = static_cast<bool>((x >> 4) & 1);
    r[4] = static_cast<bool>((x >> 3) & 1);
    r[5] = static_cast<bool>((x >> 2) & 1);
    r[6] = static_cast<bool>((x >> 1) & 1);
    r[7] = static_cast<bool>((x >> 0) & 1);

之后，看能不能找到 SIMD 指令，把 8 个位右移和 8 个位与一块儿算。

我们首先从 SSE 系列的指令集试试看！这一系列主要有 SSE SSE2 SSE3 SSSE3 SSE4.1 SSE4.2 6 代指令集，从中任找能用的指令都行。

SSE 系列指令集使用的是 128 位的 xmm 寄存器，而我们的目标是生成 8 个 bool，即 64 位，尽管浪费了一半，但是可行！

起手就是一个 _mm_set1_epi8，用于把寄存器里填上 16 个一样的 8 位变量 x：

#include <stdio.h>
#include <array>
#include <emmintrin.h>

std::array<bool, 8>
bitarray(unsigned char x)
{
    __m128i xmmvx = _mm_set1_epi8(x); // SSE2

    std::array<bool, 8> r;
    return r;
}

然后，让这 16 个 x 分别位移不同的长度就好了。

非常可惜的是，找遍了 SSE 系列的指令集，都没有一个“将矢量分别位右移不同长度”的指令。

那曲线救国，用整数除法等价替换掉位右移？

不好意思，SSE 只支持浮点数的除法，木有整数的除法。

那分别位左移不同长度，然后再统一右移相同长度？

坏消息是，分别位左移不同长度的指令同样也没有；而好消息是，位左移可用乘法等价代替。

坏消息 2 是，SSE 不支持 8 位数的乘法；而好消息 2 是，是支持 16 位的。

所以，接下来的路线，就是尝试改用 16 位整数做计算。而由于用的是 16 位整数，比我们目标—— 8 位的 bool 要宽，所以处理起来会略微有些麻烦。

我们这边改用 8 个 16 位整数，以 x = 0b11100100 为例，分语句讲解处理手法：

  1. 使用 _mm_set1_epi16 将 xmm 寄存器里填上 8 个一样的 16 位的整数，得到向量 xmmvx：
bitarray(unsigned char x)
{
    __m128i xmmvx = _mm_set1_epi16(x); // SSE2

xmmvx is
00000000'11100100
00000000'11100100
00000000'11100100
00000000'11100100
00000000'11100100
00000000'11100100
00000000'11100100
00000000'11100100
  1. 使用乘法代替位左移，将向量位移不同的长度：
    __m128i xmmvshift = _mm_set_epi16(
        1 << 7, 1 << 6, 1 << 5, 1 << 4,
        1 << 3, 1 << 2, 1 << 1, 1 << 0
    );
    xmmvx = _mm_mullo_epi16(xmmvx, xmmvshift); // SSE2

xmmvx is
00000000'11100100
00000001'11001000
00000011'10010000
00000111'00100000
00001110'01000000
00011100'10000000
00111001'00000000
01110010'00000000
  1. 将向量统一向右位移 7 个比特。2 与 3 合起来的效果，就是将向量位右移不同的长度了。
    xmmvx = _mm_srli_epi16(xmmvx, 7); // SSE2

xmmvx is
00000000'00000001
00000000'00000011
00000000'00000111
00000000'00001110
00000000'00011100
00000000'00111001
00000000'01110010
00000000'11100100
  1. 将向量统一进行按位与 1 计算。经过此步骤，最末尾的比特位被保留，其他均被清零：
    xmmvx = _mm_and_si128(xmmvx, _mm_set1_epi16(1)); // SSE2

xmmvx is
00000000'00000001
00000000'00000001
00000000'00000001
00000000'00000000
00000000'00000000
00000000'00000001
00000000'00000000
00000000'00000000
  1. 但由于此时，我们想要的结果仍位于 xmm 寄存器每 16 位整数的低 8 位上，所以需要通过 shuffle 操作，将结果排列到 xmm 寄存器的前面的低 64 位上。shuffle 操作，其实是对每个单元进行重新选择排列的过程。
    __m128i xmmvshuffle = _mm_set_epi8(
        0, 0, 0, 0, 0, 0, 0, 0,
        14, 12, 10, 8, 6, 4, 2, 0
    );
    xmmvx = _mm_shuffle_epi8(xmmvx, xmmvshuffle); // SSSE3

xmmvx is
   00000001'00000001
   00000000'00000001
   00000001'00000000
   00000000'00000000
// xmm 高位已废弃
// 00000001'00000001
// 00000001'00000001
// 00000001'00000001
// 00000001'00000001
  1. 最后一步，只需将 xmm 寄存器的低 64 位取出即可：
    std::array<bool, 8> r;
    _mm_storeu_si64(r.data(), xmmvx); // SSE2

完整代码如下，编译器的编译选项至少需支持到 SSSE3（在 gnu 或 clang 上，需添加 -mssse3 选项）：

#include <array>
#include <emmintrin.h>
#include <tmmintrin.h>

std::array<bool, 8>
bitarray(unsigned char x)
{
    __m128i xmmvx = _mm_set1_epi16(x); // SSE2
    __m128i xmmvshift = _mm_set_epi16(
        1 << 7, 1 << 6, 1 << 5, 1 << 4,
        1 << 3, 1 << 2, 1 << 1, 1 << 0
    );
    xmmvx = _mm_mullo_epi16(xmmvx, xmmvshift); // SSE2
    xmmvx = _mm_srli_epi16(xmmvx, 7); // SSE2
    xmmvx = _mm_and_si128(xmmvx, _mm_set1_epi16(1)); // SSE2
    __m128i xmmvshuffle = _mm_set_epi8(
        0, 0, 0, 0, 0, 0, 0, 0,
        14, 12, 10, 8, 6, 4, 2, 0
    );
    xmmvx = _mm_shuffle_epi8(xmmvx, xmmvshuffle); // SSSE3

    std::array<bool, 8> r;
    _mm_storeu_si64(r.data(), xmmvx); // SSE2

    return r;
}

gcc 生成的汇编是这样的（剩下没截全的几行只是常量表了），效果已经比纯 C++ 语句写出的版本好太多了。

image.png

那此处有何优化思路呢？

有的读者可能会想到，统一位右移那一步（步骤 3）完全可以省掉了呀。位左移的时候，大家都再多移一位，直接让数据保持在 xmm 寄存器每 16 位的高位上；再通过 shuffle 指令，排列到我们需要的整个寄存器的低 64 位就好了呀。

于是，代码可优化如下（注意一下，and 操作的掩码也得改成 0x100 了哦）：

#include <array>
#include <emmintrin.h>
#include <tmmintrin.h>

std::array<bool, 8>
bitarray(unsigned char x)
{
    __m128i xmmvx = _mm_set1_epi16(x); // SSE2
    __m128i xmmvshift = _mm_set_epi16(
        1 << 8, 1 << 7, 1 << 6, 1 << 5,
        1 << 4, 1 << 3, 1 << 2, 1 << 1
    );
    xmmvx = _mm_mullo_epi16(xmmvx, xmmvshift); // SSE2
    xmmvx = _mm_and_si128(xmmvx, _mm_set1_epi16(0x100)); // SSE2
    __m128i xmmvshuffle = _mm_set_epi8(
        0, 0, 0, 0, 0, 0, 0, 0,
        15, 13, 11, 9, 7, 5, 3, 1
    );
    xmmvx = _mm_shuffle_epi8(xmmvx, xmmvshuffle); // SSSE3

    std::array<bool, 8> r;
    _mm_storeu_si64(r.data(), xmmvx); // SSE2

    return r;
}

经此优化，指令数又精简了一条（省去了位右移一步）：

image.png

那能不能再给力一点？我们注意到，_mm_shuffle_epi8 指令需要 SSSE3 指令集的支持，而且整个函数中只有它一条指令拉高了对指令集的要求。能不能只在 SSE2 的支持范围内就实现整个功能？

是可以的，但是需要我们跳脱出那个纯 C++ 实现的算法思路。

这里给出新的计算思路如下，仍是以 x = 0b11100100 为例：

  1. 使用 _mm_set1_epi8 将 xmm 寄存器里填上 16 个一样的 8 位的整数，得到向量 xmmvx：
bitarray_sse2(unsigned char x)
{
	__m128i xmmvx = _mm_set1_epi8(x); // SSE2

xmmvx is
   11100100 11100100 11100100 11100100
   11100100 11100100 11100100 11100100
// xmm 高位废弃不用
// 11100100 11100100 11100100 11100100
// 11100100 11100100 11100100 11100100
  1. 和代表各个位上的掩码做位与运算：
	__m128i const xmmvandmask = _mm_set_epi8(
		0, 0, 0, 0, 0, 0, 0, 0,
		1u << 0, 1u << 1, 1u << 2, 1u << 3,
		1u << 4, 1u << 5, 1u << 6, 1u << 7
	); // SSE2
	xmmvx = _mm_and_si128(xmmvx, xmmvandmask); // SSE2

xmmvx is
   10000000 01000000 00100000 00000000
   00000000 00000100 00000000 00000000
// 00000000 00000000 00000000 00000000
// 00000000 00000000 00000000 00000000
  1. 和 1 做取 Min 运算，这步比较关键。如果先前的值各位都为 0 的话，结果就是 0，而先前但凡任意一位有 1，则运算结果为 1。从而巧妙地实现了规整化为 bool 的操作。
	xmmvx = _mm_min_epu8(xmmvx, _mm_set1_epi8(1)); // SSE2

xmmvx is
   00000001 00000001 00000001 00000000
   00000000 00000001 00000000 00000000
// 00000000 00000000 00000000 00000000
// 00000000 00000000 00000000 00000000
  1. 最后，将 xmm 寄存器的低 64 位导出即可：
	std::array<bool, 8> r;
	_mm_storeu_si64(&r[0], xmmvx); // SSE

完整代码如下：

#include <array>
#include <emmintrin.h>

std::array<bool, 8>
bitarray_sse2(unsigned char x)
{
	__m128i xmmvx = _mm_set1_epi8(x); // SSE2
	__m128i const xmmvandmask = _mm_set_epi8(
		0, 0, 0, 0, 0, 0, 0, 0,
		1u << 0, 1u << 1, 1u << 2, 1u << 3,
		1u << 4, 1u << 5, 1u << 6, 1u << 7
	); // SSE2
	xmmvx = _mm_and_si128(xmmvx, xmmvandmask); // SSE2
	xmmvx = _mm_min_epu8(xmmvx, _mm_set1_epi8(1)); // SSE2
	std::array<bool, 8> r;
	_mm_storeu_si64(&r[0], xmmvx); // SSE
	return r;
}

生成的汇编是这样的：

image.png

已经非常给力了！

能再给力一点吗？

这里介绍一下更新的 BMI 系列指令集。

BMI 和 BMI2 指令集是 intel 在 haswell 架构中新增的指令集。BM 是 bit manipulation 的缩写，正如其名，它们是专为位操作场景服务的。

我们需要用到 BMI2 中的 64 位 pdep 指令：

image.png

如果将掩码设为 0b 00000001 00000001 ... 00000001 的话，它就可以将 x 的二进制的值逐位放置在掩码置 1 的位置上。

#include <stdint.h>
#include <array>
#include <immintrin.h>

std::array<bool, 8> bitarray_bmi2(unsigned char x)
{
    uint64_t mask = 0x0101010101010101;

    uint64_t t = _pdep_u64(static_cast<uint64_t>(x), mask);
    std::array<bool, 8> r;
    *(reinterpret_cast<uint64_t*>(static_cast<void*>(&r[0]))) = t;
    return r;
}

int main()
{
    for (auto b : bitarray_bmi2(0b11100100)) {
        cout << b;
    }
    cout << endl;
}

// 输出 00100111

很可惜，和我们需要的结果是反的。

但是不要紧，我们有 bswap，一个 x86 非常早就支持的转换大小端的指令，把字节序翻转过来。

所以最终写成：

#include <stdint.h>
#include <array>
#include <immintrin.h>

std::array<bool, 8> bitarray_bmi2(unsigned char x)
{
    uint64_t mask = 0x0101010101010101;

    uint64_t t = _bswap64(_pdep_u64(static_cast<uint64_t>(x), mask));
    std::array<bool, 8> r;
    *(reinterpret_cast<uint64_t*>(static_cast<void*>(&r[0]))) = t;
    return r;
}

生成的汇编是这样的：

image.png

寥寥数语，毫无废话。

总结

如果目标环境支持 BMI2 指令集，可用以下方式求解本问题：

#include <stdint.h>
#include <array>
#include <immintrin.h>

std::array<bool, 8> bitarray_bmi2(unsigned char x)
{
    uint64_t mask = 0x0101010101010101;

    uint64_t t = _bswap64(_pdep_u64(static_cast<uint64_t>(x), mask));
    std::array<bool, 8> r;
    *(reinterpret_cast<uint64_t*>(static_cast<void*>(&r[0]))) = t;
    return r;
}

如果不支持 BMI2 但支持 SSE2，可用以下方式求解：

std::array<bool, 8>
bitarray_sse2(unsigned char x)
{
	__m128i xmmvx = _mm_set1_epi8(x); // SSE2
	__m128i const xmmvandmask = _mm_set_epi8(
		0, 0, 0, 0, 0, 0, 0, 0,
		1u << 0, 1u << 1, 1u << 2, 1u << 3,
		1u << 4, 1u << 5, 1u << 6, 1u << 7
	); // SSE2
	xmmvx = _mm_and_si128(xmmvx, xmmvandmask); // SSE2
	xmmvx = _mm_min_epu8(xmmvx, _mm_set1_epi8(1)); // SSE2
	std::array<bool, 8> r;
	_mm_storeu_si64(&r[0], xmmvx); // SSE
	return r;
}

再不支持的话，还有 MMX 版本的（和 SSE2 版思路一致，本文不再赘述）：

#include <stdint.h>
#include <array>
#include <mmintrin.h>

std::array<bool, 8>
bitarray_mmx(unsigned char x)
{
    __m64 mmvx = _mm_set1_pi8(x); // MMX
    __m64 const mmvandmask = _mm_set_pi8(
        1u << 0, 1u << 1, 1u << 2, 1u << 3,
        1u << 4, 1u << 5, 1u << 6, 1u << 7
    ); // MMX
    mmvx = _mm_and_si64(mmvx, mmvandmask); // MMX
    mmvx = _mm_cmpeq_pi8(mmvx, _mm_set1_pi8(0)); // MMX    
    mmvx = _mm_andnot_si64(mmvx, _mm_set1_pi8(1)); // MMX
    
    std::array<bool, 8> r;
    reinterpret_cast<__m64&>(r) = mmvx; // MMX
    return r;
}

这个指令集目前的 x64 机器基本都支持了，再衰的话，就用原版的循环吧。

对于 ARM 平台的支持 neon 指令集的设备，亦可使用与 SSE2 版相同的思路去做优化。

本文完

avatar
文章
阅读
粉丝