使用 SIMD 指令批量检测位于集合中的字节

与一种查找表分解法

在进行文本处理时，判断某些字节是否属于特定集合是一项 基础 操作。 对于标量实现而言，针对不同大小和特性的字节集合，判断一个字节是否存在于其中通常采用分支判断、位图或查找表等方法。 利用 SIMD 指令对这一操作进行批量处理，可以在不显著增加单次操作延迟的情况下，大幅提升吞吐量，实现远超标量实现的性能。这种批量化的处理技术又被称作 向量化分类。

向量化分类并非新事物：在受 Lemire 相关工作 启发并琢磨出一种查找表分解方法后，发现该方法与 Parsing Gigabytes of JSON per Second 中提出的方法实质上是一致的（尽管原文中的描述相对更加简洁和难以理解）。

此外，另一篇 相关的文章 也几乎完整地介绍了向量化分类的各种场景和对应的实现。幸运的是，该文给出的“特例情况”的条件过于苛刻，因此在重述算法之外，本文还能总结一些更新的理解与应用。 在接下来的部分，本文将介绍最常用的一种基于查找表的向量化分类实现以及一个特例情况下的简化实现。 对于前者，将采用一种分解思路来描述计算所需查找表的方法。

ARM64 (Neon) 和 x86-64 (SSSE3) 指令集中的表查找/字节混洗指令

TBL

在 Neon 指令集中，ARM 提供了 tbl 指令，用于执行表的向量化查找。 其中最常用的一种变体，用于单个长度为 128 位、每个元素为 1 个字节、共包含 16 个通道的表（恰好可以放置在单个向量寄存器中）。 取决于要进行批量查找的源索引寄存器长度，其指令格式为 tbl vd.16b, {vn.16b}, vm.16b 或 tbl vd.8b, {v16.8b}, vm.8b。 这条指令变体的语义官方描述如下：

此指令从索引源 SIMD 和 FP 寄存器的向量元素中读取每个值，将每个结果作为索引，在由一个源表 SIMD 和 FP 寄存器描述的字节表中进行查找，将查找结果放入向量中，并将向量写入目标 SIMD 和 FP 寄存器。如果索引超出表的范围，则该查找结果为 0。

举例而言，假设有一个 16 元素的字节查找表， 其内容为 [0xF, 0xE, 0xD, 0xC, 0xB, 0xA, 0x9, 0x8, 0x7, 0x6, 0x5, 0x4, 0x3, 0x2, 0x1, 0x0] （即索引 0 处存储的值为 15，索引 1 处存储的值为 14，以此类推，直到索引 15 处存储的值为 0）。 使用这个查找表可以将一个半字节（4 位）的值作为索引转换为 15 减去该值的结果，而如果索引超过半字节的值域（>= 16），则会产生 0 作为结果。 抛开这个查找表的实际作用不谈，如果有一个长度为 8 的字节数组，例如 [4, 2, 17, 8, 5, 1, 10, 11]， 那么通过单条 tbl 指令，即可使用上述查找表将这个数组转换为 [11, 13, 0, 7, 10, 14, 5, 4]。 在高性能 ARM64 架构实现（如 M1 系列芯片）上，无论是大核还是小核， 使用单个查找表的 tbl 指令延迟为 2 个周期，吞吐量为 0.25。

PSHUFB

尽管并没有直接用于查找表的指令，Intel 早在 SSSE3 指令集中就提供了 pshufb 指令。 当用于长度为 128 位的 xmm 寄存器时，其指令格式为 pshufb xmm1, xmm2/m128。 该指令变体的语义官方描述如下：

PSHUFB 根据源操作数（第二个操作数）中的混洗控制掩码，对目标操作数（第一个操作数）中的字节进行原地混洗。该指令对目标操作数中的数据进行重新排列，而不会影响混洗掩码。如果混洗控制掩码中每个字节的最高位（bit[7]）被置为 1，则在结果字节中写入常量零。混洗控制掩码中的每个字节形成一个索引，用于重新排列目标操作数中对应的字节。每个索引的值是混洗控制字节的最低 4 位。

正如其名“混洗”（Shuffle）所暗示的那样，其作用是使用索引源寄存器中的字节值作为目标索引， 指示目标寄存器中的字节应当如何重新排列。当索引越界时，不同于 tbl 指令总会产生 0， pshufb 只会在索引的最高位被置为 1 时产生 0。在其余情况下，它仅从字节的低 4 位中提取索引值。

虽然并非专用于表查找，pshufb 仍然能服务于该目的。 在进行查找操作时，输入的查找表会被原地更新为输出。 也就是说，这条指令里的 xmm1 寄存器同时起到了存储查找表和结果的作用， 而 xmm2 对应着输入寄存器。

查找表大小的限制

上述两条指令中，都使用了 128 位长的字节查找表，索引值为 4 位。 在 Neon 指令集中 tbl 可以最多被拓展到使用 64 字节的查找表， 但这仍然无法覆盖字节分类任务中所需的 256 种情况， 且速度无法与使用 16 字节查找表的指令相匹敌，更别说跨平台的可用性了。

在这种情形下，可以将单次查找拆分为两次查找。 这样的分解并不总是能成功：它对目标集合的特性存在要求。 然而，现实场景中的大部分具有使用价值的集合都满足这一要求。

两次查找法

分解即是升维

考虑使用一个全宽查找表进行单次查找的原理： 将待查元素作为一个 一维坐标 去获取一个对应的值， 这个值是预先计算过的，代表了待查元素的某种性质（针对本文所述的检测字节是否位于集合中的任务， 只需要 1 位的值用于表示该元素是否存在于集合中 —— 这刚好对应了位图）。

那么，是否可以通过 升维 的方式，使用 二维坐标 去更好地刻画信息， 从而缩短所需的查找表/位图长度？结论是肯定的。

举个例子：考虑识别字符集合 { '{', '}', '[', ']' } 中的字符。 它在一维的 256 位的位图中需要设置 4 位。

然而，考虑这些字符在二维 ASCII 字符表中的位置：

Brackets and braces in ASCII table

上表中的列标签代表了 ASCII 字符的低 4 位数值，并且行标签代表了 ASCII 字符的高 4 位数值。

从上表中可以发现，目标字符位于两列两行形成的 4 个交叉点上。 因此，尽管仍然需要设置 4 位（分别为 2 列和 2 行），但这次只需要两个长 16 位的查找表，分别对应列和行。 进行查找时，先将待查字节的低 4 位作为索引，从列查找表中获取一位；同时使用高 4 位作为索引，从行查找表中获取另一位。 将得到的 2 位数据进行按位与操作，即可判断该元素是否位于集合中。 这种定位方式与扫描键盘矩阵的思路颇为相似。借助这一方案，查找表的长度从 256 位缩短到了 32 位。

由于这种行列选择的特性，得到的实际上只能是原始字节矩阵的子矩阵， 这无疑 限制了能够表达的集合范围。有什么办法可以缓解这一约束吗？

额外的位带来更美妙的分解

在上述计算中使用了位图，但在实践过程中， 由于硬件指令的限制，将值映射到单个字节而非单个位的查找表实现往往更受青睐。 与之对应的，对于二维分解的查找表而言，单个查找表至少需要 16 个字节， 而这也恰好对应了可移植性最强的 128 位向量运算。

因此，在检测位于集合中的字节的任务中， 拥有一整个字节（即 8 位）可以被共同用于描述目标集合。

如果目标集合是英文字母集合，可以使用 2 位来完成识别任务。 对应的 ASCII 字符表如下所示：

ASCII table with alpha characters

在这里，使用单个比特（掩码 0b01）来表示字母 P 到 Z 以及 p 到 z。这些字母的位置是被 0-10 列和 5, 7 行所选中的子矩阵。 此外，用另一比特（掩码 0b10）来表示字母 A 到 M 与 a 到 m。 这些字母的位置是被 1-15 列和 4, 6 行所选中的子矩阵。

然后，就可以生成互相独立的查找表 —— 通过将对应位置的所有掩码进行按位或操作。 按位与操作实质上表明对两个子集进行了求并。在这个例子中，可以生成值为 [0x1, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x2, 0x2, 0x2, 0x2, 0x2] 的列查找表 和值为 [0, 0, 0, 0, 0x2, 0x1, 0x2, 0x1, 0, 0, 0, 0, 0, 0, 0, 0] 的行查找表。

当进行实际查找时，分别从列和行查找表中获取对应半字节映射到的值，然后进行按位与操作。 如果结果非 0 —— 即任意一位被设置 —— 即可说明该字节位于其中一个子集，进而属于期望的集合。

可以观察到分解并非是唯一的。比如在上述例子里，也可以消耗更多的位来分割集合：

ASCII table with alpha characters

那么，是否存在一种方法，来找到一个给定字节集合的最优分解呢？

理论上来说，这可以总结为最小化覆盖一个图形的子矩阵数量的问题。 但由于许多覆盖问题都是 NP 难的，情况可能会比较复杂。

但从工程角度考虑，无论使用三个比特还是两个比特，执行查找的总指令数是相同的。 因此只要能够在八个比特的限制下找到任意一个可行的分解，就足以提供二次查找法的最佳性能。

只需要找到超过 8 个行列都不相同的元素，即可轻易构造出二次查找法无法分解的字节集合。 不过这样的集合在实际场景中并不常见。

实现

计算完查找表后，即可将其用于 SIMD 并行检测。 以 Neon 指令集为例，以下代码可以被用来批量跳过连续的字母：

#include <arm_neon.h>
#include <stdlib.h>
#include <assert.h>

char *batch_skip_alpha(char *str, size_t len) {
  assert(len >= 16);

  // the mask for the low nibble
  uint8x16_t col_mask = (uint8x16_t) {
    0x1, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3,
    0x3, 0x3, 0x3, 0x2, 0x2, 0x2, 0x2, 0x2
  };

  // the mask for the high nibble
  uint8x16_t row_mask = (uint8x16_t) {
    0x0, 0x0, 0x0, 0x0, 0x2, 0x1, 0x2, 0x1,
    0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0
  };

  for (;;) {
    // load 16 bytes from str
    uint8x16_t src = vld1q_u8((const uint8_t *)str);

    uint8x16_t low = vqtbl1q_u8(col_mask, vandq_u8(src, vdupq_n_u8(0x0f)));
    uint8x16_t high = vqtbl1q_u8(row_mask, vshrq_n_u8(src, 4));

    // combine the lookup result for the low and high nibbles
    uint8x16_t result = vtstq_u8(low, high);

    // generate the corresponding mask
    uint8x8_t mask = vshrn_n_u16(result, 4);

    uint64_t matches = vget_lane_u64(vreinterpret_u64_u8(mask), 0);

    if (matches != (uint64_t)-1ll) {
      str += __builtin_ctzll(~matches) >> 2;
      break;
    }

    str += 16;
  }

  return str;
}

提示：代码中使用了 将 x86 向量位掩码优化移植到 ARM Neon 的通用技巧。