Parsing Protobuf at 2+GB/s: How I Learned To Love Tail Calls in C

August 31, 2024

翻译自：https://blog.reverberate.org/2021/04/21/musttail-efficient-interpreters.html

使用 `musttail` 优化解释器循环

一个令人兴奋的新特性刚刚合并到 Clang 编译器的主分支。使用 [[clang::musttail]] 或 __attribute__((musttail)) 语句属性，现在可以在 C、C++ 和 Objective-C 中获得保证的尾调用。

虽然尾调用通常与函数式编程风格相关联，但我对它们的兴趣纯粹出于性能原因。事实证明，在某些情况下，我们可以使用尾调用从编译器中获得比其他方法更好的代码 - 至少在当前的编译器技术下 - 并且无需降级到汇编语言。

将这种技术应用于 protobuf 解析已经产生了惊人的结果：我们已经成功地展示了超过 2GB/s 的 protobuf 解析速度，是之前最先进技术的两倍多。有多种技术促成了这种加速，因此“尾调用 == 2 倍加速”是一个错误的结论。但尾调用是使这种加速成为可能的关键部分。

在这篇博文中，我将描述为什么尾调用是一种如此强大的技术，我们是如何将其应用于 protobuf 解析的，以及这种技术如何推广到解释器。我认为，所有用 C 语言编写的主要语言解释器（Python、Ruby、PHP、Lua 等）都可以通过采用这种技术来获得显著的性能提升。主要缺点是可移植性：目前 musttail 是一个非标准的编译器扩展，虽然我希望它能流行起来，但要过一段时间才能广泛传播到足以让你系统的 C 编译器支持它。也就是说，在构建时，如果检测到 musttail 不可用，可以牺牲一些效率来换取可移植性。

尾调用基础

尾调用是任何位于函数尾部的调用，即在函数返回之前要执行的最后一个动作。当进行 尾调用优化 时，编译器为尾调用发出一个 jmp 指令，而非 call 指令。这跳过了通常允许被调用者 g() 返回到调用者 f() 的 bookkeeping 操作，例如创建新的堆栈帧或压入返回地址。相反，f() 直接跳转到 g()，就好像它是同一个函数的一部分一样，而 g() 直接返回到调用 f() 的任何函数。这种优化是安全的，因为一旦尾调用开始，f() 的堆栈帧就不再需要了，因为它不再可能访问 f() 的任何局部变量。

虽然这看起来像是一个普通的优化，但它具有两个非常重要的特性，为我们能够编写的算法类型打开了新的可能性。首先，它将连续进行 n 次尾调用时的堆栈内存从 O(n) 降低到 O(1)，这很重要，因为堆栈内存是有限的，而且堆栈溢出会导致程序崩溃。这意味着某些算法实际上只有在执行了这种优化后才安全写入。其次，jmp 消除了 call 的性能开销，使得函数调用与任何其它分支一样高效。这两个特性使我们能够使用尾调用作为正常迭代控制结构（如 for 或 while）的有效替代方案。

这绝不是一个新想法，实际上它至少可以追溯到 1977 年，当时 Guy Steele 发表了一篇完整的论文，论证了过程调用比 GOTO 具有更简洁的设计，而尾调用优化可以使它们一样快。这是 “Lambda Papers” 之一，这些论文是在 1975 年到 1980 年间撰写的，期间发展了许多 Lisp 和 Scheme 的基础思想。

尾调用优化甚至对 Clang 来说也不是什么新鲜事：像 GCC 和许多其它编译器一样，Clang 已经能够优化尾调用。实际上，我们第一个示例中的 musttail 属性根本没有改变编译器的输出：Clang 已经可以在 -O2 下优化尾调用。

新的是 保证性。虽然编译器通常能够成功地优化尾调用，但这只是尽力而为，而不是可以依赖的东西。特别是，这种优化很可能不会在非优化构建中发生：

此处，尾调用被编译成真正的 call，所以我们又回到了 O(n) 堆栈空间。这就是我们需要 musttail 的原因：除非我们能从编译器那里得到保证，确保我们的尾调用将在所有构建模式下始终被优化，否则使用尾调用进行迭代的算法是不安全的。如果代码只有在启用优化时才能正常工作，那将是一个非常严重的限制。

解释器循环的麻烦

编译器是令人难以置信的技术，但它们并不完美。LuaJIT 的作者 Mike Pall 决定用汇编语言而不是 C 语言编写 LuaJIT 2.x 的解释器，他将此决定视为解释为什么 LuaJIT 的解释器如此快的主要因素。他后来详细介绍了为什么 C 编译器难以处理解释器主循环。他最核心的两点是：

函数越大，其控制流越复杂和相互关联，编译器的寄存器分配器就越难将最重要的数据保存在寄存器中。
当快速路径和慢速路径混合在同一个函数中时，慢速路径的存在会影响快速路径的代码质量。

这些观察结果与我们在优化 protobuf 解析时的经验非常吻合。好消息是尾调用可以帮助解决这两个问题。

将解释器循环与 protobuf 解析器进行比较可能看起来很奇怪，但 protobuf 的线格式的性质使它们比你想象的更相似。protobuf 的线格式是一系列标签/值对，其中标签包含字段编号和线类型。这个标签类似于解释器的操作码：它告诉我们需要执行什么操作来解析该字段的数据。与解释器操作码一样，protobuf 字段编号可以以任何顺序出现，因此我们必须准备好随时调度到代码的任何部分。

编写这种解析器的自然方法是让一个 while 循环围绕一个 switch 语句，实际上，这基本上是 protobuf 解析的最新技术，从 protobuf 存在的那一刻起就是如此。例如，以下是一些来自 protobuf 当前 C++ 版本的解析代码。如果我们以图形方式表示控制流，我们会得到类似于以下内容：

但这是不完整的，因为几乎在每个阶段都可能出现问题。线类型可能错误，或者我们可能看到一些损坏的数据，或者我们可能只是碰巧到达了当前缓冲区的末尾。所以完整的控制流图看起来更像这样。

我们希望尽可能地停留在快速路径（蓝色）上，但当我们遇到一个困难的情况时，我们必须执行一些回退代码来处理它。这些回退路径通常比快速路径更大、更复杂，涉及更多数据，并且通常甚至进行对其他函数的越界调用来处理更复杂的情况。

理论上，这种控制流图与配置文件相结合，应该能为编译器提供生成最优代码所需的所有信息。在实践中，当一个函数如此庞大且相互关联时，我们经常发现自己与编译器作斗争。在我们需要它将重要变量保存在寄存器中时，它却溢出了那个变量。它提升了我们想缩放到回退函数调用的堆栈帧操作。它合并了我们希望为了分支预测原因而保持分离的相同代码路径。这种体验最终可能会让人感觉像戴着手套弹钢琴一样。

使用尾调用改进解释器循环

上面的分析主要只是对 Mike 关于解释器主循环的观察的复述。但是，我们并没有像 Mike 使用 LuaJIT 2.x 那样降低到汇编语言，而是发现面向尾调用的设计可以为我们提供从 C 语言获得几乎最佳代码所需的控制。我和我的同事 Gerben Stavenga 一起研究了这个问题，他提出了大部分的设计。我们的方法类似于 wasm3 WebAssembly 解释器的设计，该解释器将这种模式描述为一个 “元机器”。

我们 2+GB/s protobuf 解析器的代码提交给了 upb，这是一个用 C 语言编写的小型 protobuf 库，在 pull/310 中提交。虽然它完全可以工作，并且通过了所有 protobuf 兼容性测试，但它还没有在任何地方推出，并且该设计也没有在 protobuf 的 C++ 版本中实现。但现在 musttail 在 Clang 中可用（并且 upb 已更新为使用它），完全生产化快速解析器的最大障碍之一已经消除。

我们的设计不再使用单个大型解析函数，而是为每个操作提供一个自己的小型函数。每个函数都尾调用下一个操作序列。例如，以下是一个用于解析单个定长字段的函数。（这段代码是从 upb 中的实际代码简化而来的；有很多设计细节我在这篇文章中没有提到，但希望在未来的文章中涵盖）。

#include <stdint.h>
#include <stddef.h>
#include <string.h>

typedef void *upb_msg;
struct upb_decstate;
typedef struct upb_decstate upb_decstate;

// 传递给每个解析函数的标准参数集。
// 由于 x86-64 的调用约定，这些参数将通过寄存器传递。
#define UPB_PARSE_PARAMS                                          \
  upb_decstate *d, const char *ptr, upb_msg *msg, intptr_t table, \
      uint64_t hasbits, uint64_t data
#define UPB_PARSE_ARGS d, ptr, msg, table, hasbits, data

#define UNLIKELY(x) __builtin_expect(x, 0)
#define MUSTTAIL __attribute__((musttail))

const char *fallback(UPB_PARSE_PARAMS);
const char *dispatch(UPB_PARSE_PARAMS);

// 解析使用 1 字节标签（字段 1-15）的 4 字节定长字段的代码
const char *upb_pf32_1bt(UPB_PARSE_PARAMS) {
  // 解码“data”，其中包含关于该字段的信息
  uint8_t hasbit_index = data >> 24;
  size_t ofs = data >> 48;

  if (UNLIKELY(data & 0xff)) {
    // 线类型不匹配（dispatch 函数将预期的线类型与实际的线类型进行异或运算，因此 data & 0xff == 0 表示匹配）
    MUSTTAIL return fallback(UPB_PARSE_ARGS);
  }

  ptr += 1;  // 跳过标签

  // 将数据存储到消息中
  hasbits |= 1ull << hasbit_index;
  memcpy((char*)msg + ofs, ptr, 4);

  ptr += 4;  // 跳过数据

  // 调用 dispatch 函数，该函数将读取下一个标签，并分支到正确的字段解析函数。
  MUSTTAIL return dispatch(UPB_PARSE_ARGS);
}

对于这样一个小型而简单的函数，Clang 为我们提供了几乎无法超越的代码。

upb_pf32_1bt:                           # @upb_pf32_1bt
        mov     rax, r9
        shr     rax, 24
        bts     r8, rax
        test    r9b, r9b
        jne     .LBB0_1
        mov     r10, r9
        shr     r10, 48
        mov     eax, dword ptr [rsi + 1]
        mov     dword ptr [rdx + r10], eax
        add     rsi, 5
        jmp     dispatch                        # TAILCALL
.LBB0_1:
        jmp     fallback                        # TAILCALL

请注意，没有序言或尾声指令，没有寄存器溢出，实际上根本没有使用堆栈。唯一的出口是来自两个尾调用的 jmp，但不需要任何代码来转发参数，因为参数已经位于正确的寄存器中。我们几乎唯一希望的改进是获得尾调用的条件跳转，jne fallback，而不是 jne 后面跟着 jmp。

译注：
序言（prologue）指令位于函数体的开头，用于保存上下文并创建栈帧，例如：
push ebp ; 保存调用者的栈帧指针
mov ebp, esp ; 将当前栈指针设置为新的栈帧指针
sub esp, N ; 为局部变量保留N字节的栈空间
push ebx ; 保存需要使用的寄存器（比如ebx）
push esi
push edi
尾声（epilogue）指令位于函数体的末尾，用于清理和恢复在 prologue 中保存的状态，例如：
pop edi       ; 恢复寄存器
pop esi
pop ebx
mov esp, ebp  ; 恢复栈指针
pop ebp       ; 恢复栈帧指针
ret           ; 返回调用点
另见：X86调用约定

如果你查看这段代码的汇编指令，而没有符号信息，你将没有理由知道这是一个完整的函数。它也可能是来自更大函数的基本块。而这本质上正是我们正在做的事情。我们正在将一个在概念上是一个大型复杂函数的解释器循环分解成块，并通过尾调用将控制流从一个块转移到另一个块。我们在每个块边界（至少对于六个寄存器）都完全控制寄存器分配，并且通过对每个函数使用相同的一组参数，我们消除了在一次调用到下一次调用之间来回移动任何值的必要性。只要函数足够简单，不会溢出这六个寄存器，我们就实现了将最重要的状态始终保存在所有快速路径中的寄存器中的目标。

我们可以独立地优化每个指令序列，而且至关重要的是，编译器也会将每个序列视为独立的，因为它们位于不同的函数中（如果需要，我们可以使用 noinline 阻止内联）。这解决了我们之前描述的问题，即回退路径中的代码会降低快速路径的代码质量。如果我们将慢速路径放在与快速路径完全不同的函数中，我们可以保证快速路径不会受到影响。我们上面看到的漂亮的汇编序列实际上是冻结的，不受我们在解析器其他部分所做的任何更改的影响。

如果我们将这种模式应用于 Mike 来自 LuaJIT 的示例，我们或多或少可以用只有轻微代码质量缺陷的手写汇编指令来匹配他的代码：

#define PARAMS unsigned RA, void *table, unsigned inst, \
               int *op_p, double *consts, double *regs
#define ARGS RA, table, inst, op_p, consts, regs
typedef void (*op_func)(PARAMS);
void fallback(PARAMS);

#define UNLIKELY(x) __builtin_expect(x, 0)
#define MUSTTAIL __attribute__((musttail))

void ADDVN(PARAMS) {
    op_func *op_table = table;
    unsigned RC = inst & 0xff;
    unsigned RB = (inst >> 8) & 0xff;
    unsigned type;
    memcpy(&type, (char*)&regs[RB] + 4, 4);
    if (UNLIKELY(type > -13)) {
        return fallback(ARGS);
    }
    regs[RA] += consts[RC];
    inst = *op_p++;
    unsigned op = inst & 0xff;
    RA = (inst >> 8) & 0xff;
    inst >>= 16;
    MUSTTAIL return op_table[op](ARGS);
}

生成的汇编指令是：

ADDVN:                                  # @ADDVN
        movzx   eax, dh
        cmp     dword ptr [r9 + 8*rax + 4], -12
        jae     .LBB0_1
        movzx   eax, dl
        movsd   xmm0, qword ptr [r8 + 8*rax]    # xmm0 = mem[0],zero
        mov     eax, edi
        addsd   xmm0, qword ptr [r9 + 8*rax]
        movsd   qword ptr [r9 + 8*rax], xmm0
        mov     edx, dword ptr [rcx]
        add     rcx, 4
        movzx   eax, dl
        movzx   edi, dh
        shr     edx, 16
        mov     rax, qword ptr [rsi + 8*rax]
        jmp     rax                             # TAILCALL
.LBB0_1:
        jmp     fallback

我在这里看到的唯一改进机会，除了前面提到的 jne fallback 问题之外，是编译器出于某种原因不想生成 jmp qword ptr [rsi + 8*rax]。相反，它更喜欢加载到 rax 中，然后跟随 jmp rax。这些是轻微的代码生成问题，希望可以在 Clang 中通过并不太大的工作量来解决。

限制

这种方法最大的问题之一是，如果存在任何非尾调用，这些漂亮的汇编序列就会被灾难性地悲观化。任何非尾调用都会导致创建堆栈帧，并且大量数据会溢出到堆栈中。

#define PARAMS unsigned RA, void *table, unsigned inst, \
               int *op_p, double *consts, double *regs
#define ARGS RA, table, inst, op_p, consts, regs
typedef void (*op_func)(PARAMS);
void fallback(PARAMS);

#define UNLIKELY(x) __builtin_expect(x, 0)
#define MUSTTAIL __attribute__((musttail))

void ADDVN(PARAMS) {
    op_func *op_table = table;
    unsigned RC = inst & 0xff;
    unsigned RB = (inst >> 8) & 0xff;
    unsigned type;
    memcpy(&type, (char*)&regs[RB] + 4, 4);
    if (UNLIKELY(type > -13)) {
        // 当我们省略 "return" 时，事情会变得非常糟糕。
        fallback(ARGS);
    }
    regs[RA] += consts[RC];
    inst = *op_p++;
    unsigned op = inst & 0xff;
    RA = (inst >> 8) & 0xff;
    inst >>= 16;
    MUSTTAIL return op_table[op](ARGS);
}

这导致了非常不幸的结果：

ADDVN:                                  # @ADDVN
        push    rbp
        push    r15
        push    r14
        push    r13
        push    r12
        push    rbx
        push    rax
        mov     r15, r9
        mov     r14, r8
        mov     rbx, rcx
        mov     r12, rsi
        mov     ebp, edi
        movzx   eax, dh
        cmp     dword ptr [r9 + 8*rax + 4], -12
        jae     .LBB0_1
.LBB0_2:
        movzx   eax, dl
        movsd   xmm0, qword ptr [r14 + 8*rax]   # xmm0 = mem[0],zero
        mov     eax, ebp
        addsd   xmm0, qword ptr [r15 + 8*rax]
        movsd   qword ptr [r15 + 8*rax], xmm0
        mov     edx, dword ptr [rbx]
        add     rbx, 4
        movzx   eax, dl
        movzx   edi, dh
        shr     edx, 16
        mov     rax, qword ptr [r12 + 8*rax]
        mov     rsi, r12
        mov     rcx, rbx
        mov     r8, r14
        mov     r9, r15
        add     rsp, 8
        pop     rbx
        pop     r12
        pop     r13
        pop     r14
        pop     r15
        pop     rbp
        jmp     rax                             # TAILCALL
.LBB0_1:
        mov     edi, ebp
        mov     rsi, r12
        mov     r13d, edx
        mov     rcx, rbx
        mov     r8, r14
        mov     r9, r15
        call    fallback
        mov     edx, r13d
        jmp     .LBB0_2

为了避免这种情况，我们尝试遵循一个规则，即只通过内联或尾调用来调用其他函数。如果一个操作有多个点可能发生不寻常的情况，而这些情况不是错误，这可能会变得很烦人。例如，当我们解析 protobuf 时，快速而常见的情况是 varint 只有一个字节长，但更长的 varint 不是错误。如果回退代码太复杂，那么在内联中处理不寻常的情况可能会影响快速路径的质量。但尾调用到回退函数没有简单的方法在处理完不寻常情况后恢复操作，因此回退函数必须能够向前推进并完成操作。这会导致代码重复和复杂性。

理想情况下，这个问题可以通过在回退函数中添加 __attribute__((preserve_most))，然后正常调用它们，而不进行尾调用来解决。preserve_most 属性使被调用者负责几乎所有寄存器的保存，这将寄存器溢出的成本转移到我们想要的回退函数中。我们对此属性进行了一些实验，但遇到了一些神秘的问题，我们无法弄清楚是什么原因造成的。这可能是我们的错误；重新审视这一点是未来的工作。[更新：2023-03-14: 事实上，这是一个 Clang 中的 bug，两个月前已修复]

另一个主要限制是 musttail 不可移植。我非常希望这个属性能流行起来，传播到 GCC、Visual C++ 和其他流行的编译器，甚至有一天能被标准化。但那一天还很遥远，那么在此期间该怎么办？

当 musttail 不可用时，我们需要为每个概念上的循环迭代执行至少一个真正的 return，而不进行尾调用。我们还没有在 upb 中实现这个回退，但我预计它将涉及一个宏，该宏根据 musttail 的可用性，要么尾调用到 dispatch，要么只是返回。