使用 SIMD 将 ascii 字符串位打包为 7 位二进制 blob

Question

相关：bitpack ascii string into 7-bit binary blob using ARM-v8 Neon SIMD - 同样的问题专门针对 AArch64 内在函数。这个问题涵盖了可移植的 C 和 x86-64 内在函数。

---

我想将 char 字符串编码为 7 位 blob，以减少 12.5% 的内存。我想尽可能快地完成它，即在编码大字符串时以最小的延迟。

这是该算法的简单实现：

void ascii_pack(const char* ascii, size_t len, uint8_t* bin) {
  uint64_t val;
  const char* end = ascii + len;

  while (ascii + 8 <= end) {
    memcpy(&val, ascii, 8);
    uint64_t dest = (val & 0xFF);

    // Compiler will perform loop unrolling
    for (unsigned i = 1; i <= 7; ++i) {
      val >>= 1;
      dest |= (val & (0x7FUL << 7 * i));
    }
    memcpy(bin, &dest, 7);
    bin += 7;
    ascii += 8;
  }

  // epilog - we do not pack since we have less than 8 bytes.
  while (ascii < end) {
    *bin++ = *ascii++;
  }
}

现在，我想用 SIMD 来加速它。我使用了下面的 SSE2 算法。我的问题：

1. 是否可以优化顺序内部循环？
2. 在大字符串上运行时它会提高吞吐量吗？

// The algo - do in parallel what ascii_pack does on two uint64_t integers
void ascii_pack_simd(const char* ascii, size_t len, uint8_t* bin) {
  __m128i val;

  __m128i mask = _mm_set1_epi64x(0x7FU);  // two uint64_t masks

  // I leave out 16 bytes in addition to 16 that we load in the loop
  // because we store into "bin" full 16 bytes instead of 14. To prevent out of bound
  // writes we finish one iteration earlier.
  const char* end = ascii + len - 32;
  while (ascii <= end) {
    val = _mm_loadu_si128(reinterpret_cast<const __m128i*>(ascii));
    __m128i dest = _mm_and_si128(val, mask);

    // Compiler unrolls it
    for (unsigned i = 1; i <= 7; ++i) {
      val = _mm_srli_epi64(val, 1);                          // shift right both integers
      __m128i shmask = _mm_slli_epi64(mask, 7 * i);    // mask both
      dest = _mm_or_si128(dest, _mm_and_si128(val, shmask));  // add another 7bit part.
    }

    // dest contains two 7 byte blobs. Lets copy them to bin.
    _mm_storeu_si128(reinterpret_cast<__m128i*>(bin), dest);
    memmove(bin + 7, bin + 8, 7);
    bin += 14;
    ascii += 16;
  }

  end += 32;  // Bring back end.
  DCHECK(ascii < end);
  ascii_pack(ascii, end - ascii, bin);
}

Answer 1

标量技巧（不需要PEXT）可以这样实现：

\n

uint64_t compress8x7bit(uint64_t x)\n{\n    x = ((x & 0x7F007F007F007F00) >> 1) | (x & 0x007F007F007F007F);\n    x = ((x & 0x3FFF00003FFF0000) >> 2) | (x & 0x00003FFF00003FFF);\n    x = ((x & 0x0FFFFFFF00000000) >> 4) | (x & 0x000000000FFFFFFF);\n    return x;\n}\n

\n

这里的想法是将相邻对连接在一起，首先将 7 位元素连接成 14 位元素，然后将它们连接成 28 位元素，最后将它们连接成一个 56 位块（这就是结果）。

\n

使用 SSSE3，您可以使用pshufb也可以连接其中两个 56 位部分（在存储它们之前）。

\n

SSE2（和 AVX2）可以做与具有 64 位元素的标量代码相同的事情，但这种方法没有利用特殊操作可能实现的任何技术（SSE2+ 有很多，每个版本都有更多），除了在 SIMD 中实现标量技巧之外，可能还有更好的事情要做。

\n

例如只是为了扔一些狂野的东西，gf2p8affineqb(0x8040201008040201, x)会将所有“丢弃”的位放在一个位置（即结果的顶部字节），并从我们想要保留的位中生成一个可靠的 56 位块。但这些位确实以奇怪的顺序结束（第一个字节将包含位 56、48、40、32、24、16、8、0，按该顺序，首先列出最低有效位）。

\n

这个顺序虽然很奇怪，但可以使用pshufb反转字节轻松解压（您也可以使用它来插入两个零），然后gf2p8affineqb(0x0102040810204080, reversedBytes)将这些位重新整理回原始顺序。

\n

下面是如何与实际 AVX2+GFNI 内在函数一起工作的草图。我不想在这里处理最后的额外部分，只是处理“主”循环，因此输入文本最好是 32 字节的倍数。适用于我的电脑 \xe2\x9c\x94\xef\xb8\x8f

\n

void compress8x7bit(const char* ascii, size_t len, uint8_t* bin)\n{\n    const char* end = ascii + len;\n    while (ascii + 31 < end) {\n        __m256i text = _mm256_loadu_si256((__m256i*)ascii);\n        __m256i transposed = _mm256_gf2p8affine_epi64_epi8(_mm256_set1_epi64x(0x8040201008040201), text, 0);\n        __m256i compressed = _mm256_shuffle_epi8(transposed, \n            _mm256_set_epi8(-1, -1, 14, 13, 12, 11, 10, 9, 8, 6, 5, 4, 3, 2, 1, 0,\n                            -1, -1, 14, 13, 12, 11, 10, 9, 8, 6, 5, 4, 3, 2, 1, 0));\n        _mm_storeu_si128((__m128i*)bin, _mm256_castsi256_si128(compressed));\n        _mm_storeu_si128((__m128i*)(bin + 14), _mm256_extracti128_si256(compressed, 1));\n        bin += 28;\n        ascii += 32;\n    }\n}\n\nvoid uncompress8x7bit(char* ascii, size_t len, const uint8_t* bin)\n{\n    const char* end = ascii + len;\n    while (ascii + 31 < end) {\n        __m256i raw = _mm256_inserti128_si256(_mm256_castsi128_si256(_mm_loadu_si128((__m128i*)bin)), _mm_loadu_si128((__m128i*)(bin + 14)), 1);\n        __m256i rev_with_zeroes = _mm256_shuffle_epi8(raw, \n            _mm256_set_epi8(7, 8, 9, 10, 11, 12, 13, -1, 0, 1, 2, 3, 4, 5, 6, -1,\n                            7, 8, 9, 10, 11, 12, 13, -1, 0, 1, 2, 3, 4, 5, 6, -1));\n        __m256i decompressed = _mm256_gf2p8affine_epi64_epi8(_mm256_set1_epi64x(0x0102040810204080), rev_with_zeroes, 0);\n        _mm256_storeu_si256((__m256i*)ascii, decompressed);\n        bin += 28;\n        ascii += 32;\n    }\n}\n

\n

也许有一个比在压缩器中使用两个 128 位存储和在解压缩器中使用两个 128 位加载更好的解决方案。对于 AVX512，这很容易，因为它具有全寄存器字节粒度排列，但 AVX2 具有vpshufb，它无法在构成 256 位向量的两个 128 位半部分之间移动字节。解压缩器可以执行一个有趣的加载，在它想要的数据开始之前 2 个字节开始，如下所示：_mm256_loadu_si256((__m256i*)(bin - 2))以及稍微不同的随机向量），代价是必须避免使用任一填充时出现潜在的越界错误或特殊的第一次迭代，但压缩器不能（不便宜）使用像提前 2 个字节开始的存储那样的技巧（这会破坏结果的两个字节）。

\n

顺便说一句，我这里有一些测试代码，您可以使用它们来验证您的位压缩函数是否做了正确的事情（好吧 - 只要该函数是位排列，其中某些位可能会被归零，这就可以工作作为检查，但这通常不会检测到所有可能的错误）：

\n

uint64_t bitindex[7];\nbitindex[6] = compress8x7bit(0xFFFFFFFFFFFFFFFF);\nbitindex[5] = compress8x7bit(0xFFFFFFFF00000000);\nbitindex[4] = compress8x7bit(0xFFFF0000FFFF0000);\nbitindex[3] = compress8x7bit(0xFF00FF00FF00FF00);\nbitindex[2] = compress8x7bit(0xF0F0F0F0F0F0F0F0);\nbitindex[1] = compress8x7bit(0xCCCCCCCCCCCCCCCC);\nbitindex[0] = compress8x7bit(0xAAAAAAAAAAAAAAAA);\n\nfor (size_t i = 0; i < 64; i++)\n{\n    if (i != 0)\n        std::cout << ", ";\n    if (bitindex[6] & (1uLL << i))\n    {\n        int index = 0;\n        for (size_t j = 0; j < 6; j++)\n        {\n            if (bitindex[j] & (1uLL << i))\n                index |= 1 << j;\n        }\n        std::cout << index;\n    }\n    else\n        std::cout << "_";\n}\nstd::cout << "\\n";\n

\n