使用 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 减去该值的结果，而如果索引超过半字节的值域（>= 15），则会产生 0 作为结果。 抛开这个查找表的实际作用不谈，如果我们有一个长度为 8 的字节数组，例如 [4, 2, 17, 8, 5, 1, 10, 11]， 那么通过单条 tbl 指令，我们即可使用上述查找表将这个数组转换为 [11, 13, 8, 7, 10, 14, 5, 4]。 在高性能 ARM64 架构实现，例如 M1 系列芯片上，无论是大核还是小核， 使用单个查找表的 tbl 指令延迟为 2 个周期，吞吐量为 0.25， 后者表示最多同时运行 4 条尚未完成的 tbl 指令。

PSHUFB

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

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

正如它的名字“混洗” (Shuffle) 所暗示的那样，它的作用是使用 xmm2 寄存器中位于位置 n 的字节值 m 作为目标索引， 表明 xmm1 寄存器中位置 n 的字节应该被放置于 xmm1 寄存器中位置 m 处。当索引越界时，不同于 tbl 指令总会产生 0， pshufb 只会在索引的最高位被置为 1 时，才会产生 0。 在余下的情况，它只是静默的从字节的低 4 位中提取索引值，而不管高 3 位的值具体是什么。

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

查找表大小的限制

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

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

两次查找法

分解即是升维

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

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

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

然而，考虑他们出现在 2 维 ASCII 字符表中的位置：

Brackets and braces in ASCII table

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

从上表中，可以发现，我们希望识别的字符位于两列两行形成的 4 个交叉点上。 因此，尽管仍然需要设置 4 位 - 分别为 2 列和 2 行 - 但这次我们只需要两个、 单个长 16 位的查找表，分别对应着列和行。 进行查找时，我们先将待查字节的低 4 位作为索引，从列查找表中获取一位， 同时使用待查字符的高 4 位作为索引，从行查找表中获取另一位。 将得到的 2 位数据进行按位与操作，即可获得可以判断该元素是否位于集合中的 1 位数据。 这种定位方式与扫描键盘上某个按键是否被按下的方式颇为相似， 借助这一思路，我们成功将查找表的长度从 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;
}

提示：我们使用了 使用 ARM Neon 进行位操作：超越 SSE movemask 中的技巧来生成掩码。