ESP32-S3 用 SIMD 加速 FAST 特征检测，吞吐量从 5.1MP/s 升至 11.2MP/s

我为 ESP32-S3 编写了一个利用其 128 位 SIMD 指令的加速版 FAST 特征检测器。它的性能约为参考实现的两倍，处理一帧 QVGA (320x240) 图像大约需要 ~6ms。你可以在这里找到源代码。

就这个价格而言，ESP32-S3 堪称一颗性能强大的微控制器。其不起眼的塑料封装内，置有一枚最高频率达 240MHz 的双核 CPU，以及众多外设，包括 WiFi 和 Bluetooth Low Energy 无线功能。在翻阅其技术参考手册时，我发现这款芯片支持一组有限的 SIMD 指令。对于成本比一杯普通咖啡还低的芯片来说，这相当酷。让我们看看能用它做些什么有趣的事情！

在编写 SIMD 加速版 FAST 角点检测器的过程中，我学到了 ESP32-S3 汇编的基础知识，绕开了它那些古怪的限制，最后还自己编写了一个基础寄存器分配器，basm。

又在网上浏览了一阵后，我意识到这些指令并不能算作真正的秘密，只是相对不为人知。我想亲自试验这些指令，并编写出比网上能找到的代码更高级的实现。由于最近我对计算机视觉非常感兴趣，我决定为 ESP32-S3 编写一个采用 SIMD 加速的FAST 特征检测器实现。

我先寻找使用 SIMD 指令的示例代码。在撰写本文时，我只找到了三个来源：

Larry Bank 的博客，他在其中探索了这一特性，并提供了一个最小示例
Espressif 的esp-dsp库，用于类 DSP 功能（卷积、FFT 等）
Espressif 的esp-dl库，用于片上 AI 加速。

逐一浏览完这些内容后，我如释重负，因为 ESP32-S3 上的汇编其实也没那么难。我决定采用以下总体结构：

    .section .rodata
	# read-only data here

    .section .text
    .align 4

    .global add
    .type   add,@function
# a simple function that adds two numbers
# int32_t add(int32_t a, int32_t b)
add:
    entry           sp, 48 # sp (stack pointer) is an alias for a1
    
    #define r_a a2
    #define r_b a3

	# code here
    
	add a2, r_a, r_b # a2 = a + b

    retw.n # return, with a2 containing the return value.

接下来，我开始编写 SIMD 加速版 FAST 角点预测试。本质上，它会在上下左右每个方向上采样一个像素，并检查是否存在至少三个“极端”像素。这里，极端像素是指其与中心像素的绝对差值高于某个阈值的像素。

我决定编写一个每次计算十六个像素极值计数的函数，因为每个向量寄存器可以存储十六个八位值。实现方式是在每个正交方向上采样四“块”像素，并将其与各自对应的中心像素进行比较。

然而，我在这里遇到了第一个限制——ISA 不允许直接进行非对齐读取。在假定中心像素是对齐的情况下，东侧和西侧数据块会处于非对齐状态：

幸运的是，这颗芯片有一条指令，可以将两个拼接寄存器构成的序列中的一段切片提取到另一个寄存器中。利用这一功能，我可以通过读取并提取中心数据块两侧相邻数据块中的切片，提取出东侧和西侧数据块。

好了，现在我已经有了像素数据。该进行一些比较了，对吧？并非如此。事实证明，ESP32-S3 只实现了有符号八位数的比较。由于我的所有像素数据都以无符号八位数存储，这让处理变得稍微困难了一些。如果只是将无符号数据重新解释为有符号数据，就会得到一些奇怪的结果：

>>> np.int8(np.uint8(range(0, 256, 8)))
array([   0,    8,   16,   24,   32,   40,   48,   56,   64,   72,   80,
         88,   96,  104,  112,  120, -128, -120, -112, -104,  -96,  -88,
        -80,  -72,  -64,  -56,  -48,  -40,  -32,  -24,  -16,   -8],
      dtype=int8)

[128 255] 范围会以非线性方式映射到 [-128 -1]，从而破坏所有比较运算。为了让这种方法奏效，我需要将值的范围从 [0, 255] 移到 [-128 127]。这本应是简单的减法运算，但指令集再次从中作梗：在 ESP32-S3 上，所有 SIMD 算术运算都是饱和运算，这意味着所有大于 127 的像素值最终都会被钳制在 -128，而不是下溢到各自适当的值。为了克服这个问题，我必须设法在不调用减法指令的情况下，并行地从每个数中减去 128。起初这对我来说似乎是一项艰巨的任务，但我最终找到了解决方案。

假设 x 是被减数。以下是它各个数位的位值：

-128 64 32 16 8 4 2 1

注意到什么有趣的地方了吗？减法等同于加上 -128，因此要计算 x - 128，只需给 x 中位值为 -128 的最高数位加上 1，即可让这个数“递增” -128。只需翻转 x 的最高位即可完成这一计算。由于将任意一位与 1 进行 XOR 运算等同于对该位取按位非，我们可以用位掩码 0b10000000 对 x 的最高位进行 XOR 运算，从而有选择地翻转该位，而该掩码等于 0x80。因此，x - 128 == x ^ 0x80。

现在，每当载入无符号八位值时，都会将其与 0x80 进行 XOR 运算，以把它们转换到合适的（线性）范围，使算术运算和比较再次能够正确进行。

借助这两个技巧，我得以绕过 ESP32-S3 的限制，并编写了 SIMD 加速的 FAST 角点预测试。我还使用一组类似的操作编写了一个 SIMD 加速的评分函数，不过这里不会讨论它。

最终，在我的测试中，我成功将 FAST 特征检测器的吞吐量提升了约 220%，从 5.1MP/s 提升至 11.2MP/s。这完全处于实时计算机视觉任务可接受的性能范围内，使 ESP32-S3 能够轻松处理 30fps 的 VGA 视频流。只花 $2 还不错！