Python中新JIT功能介绍

2023 年 12 月下旬（准确地说是圣诞节），CPython 核心开发人员Brandt Bucher向 Python 3.13 分支提交了一个添加 JIT 编译器的小请求。

这一更改一旦被接受，将是自 Python 3.11（也来自 Brandt）中添加的专用自适应解释器以来对 CPython 解释器的最大更改之一。

在这篇博文中，我们将了解这个 JIT，它是什么、它如何工作以及有什么好处。

什么是 JIT？
JIT（或“Just in Time”）是一种编译设计，意味着在代码第一次运行时按需进行编译。这是一个非常广泛的术语，可能意味着很多事情。我想，从技术上讲，Python 编译器已经是 JIT，因为它从 Python 代码编译为字节码。

当人们说 JIT 编译器时，他们通常指的是发出机器代码的编译器。这与 AOT（提前）编译器形成对比，例如 GNU C 编译器、GCC 或 Rust 编译器 rustc，后者生成机器代码一次并作为二进制可执行文件分发。

当你运行Python代码时，它首先被编译成字节码。网上有很多关于这个过程的演讲和视频，所以我不想重复太多，但关于 Python 字节码需要注意的重要一点是：

  • 它们对 CPU 没有任何意义，需要特殊的字节码解释器循环来执行
  • 它们是高级别的，相当于 1000 条机器指令
  • 它们与类型无关
  • 它们是跨平台的

对于一个非常简单的 Python 函数 f()，它定义了一个变量 a 并赋值为 1：

def f():
   a = 1

它编译成 4 个字节码指令，运行 dis.dis 即可看到：

>>> import dis
>>> dis.dis(f)
  2           0 LOAD_CONST               1 (1)
              2 STORE_FAST               0 (a)
              4 LOAD_CONST               0 (None)
              6 RETURN_VALUE

如果你想尝试更复杂的东西，还有一个交互性更强的反汇编器，名为 dissy。

对于这个函数，Python 编译成了 LOAD_CONST、STORE_FAST、LOAD_CONST 和 RETURN_VALUE 指令。当用 C 编写的大循环运行该函数时，这些指令将被解释。

如果要在 Python 中编写一个非常粗糙的 Python 评估循环，与 C 语言中的循环相当，它将看起来像这样：

def interpret(bytecodes):
    stack = []
    variables = {}
    for bytecode, arg in bytecodes:
        if bytecode == "LOAD_CONST":
            stack.append(arg)
        elif bytecode == "LOAD_FAST":
            stack.append(variables[arg])
        elif bytecode == "STORE_FAST":
            variables[arg] = stack.pop()
        elif bytecode == "RETURN_VALUE":
            return stack.pop()


func = (
    ("LOAD_CONST", 1),
    ("STORE_FAST", 'a'),
    ("LOAD_CONST", None),
    ("RETURN_VALUE", None)
)

如果你把测试函数的字节码交给这个解释器，它就会执行这些字节码并打印结果：
print(interpret(function))

这个带有大段 switch/if-else 语句的循环相当于 CPython 解释器循环的工作方式，尽管是简化版。CPython 由 C 语言编写，并由 C 编译器编译。为简单起见，我将用 Python 构建这个示例。

对于我们的解释器来说，每次要运行函数 func 时，它都要对每条指令进行循环，并将字节码名称（称为操作码）与每个 if 语句进行比较。这种比较和循环本身都会增加执行的开销。如果运行函数 10,000 次，而字节码从未改变（因为它们是不可变的），那么这种开销就显得多余了。与其每次调用函数时都评估循环，不如按顺序生成代码更有效率。

这就是 JIT 的作用。JIT 编译器有多种类型。Numba 就是一个 JIT。PyPy 有一个 JIT。Java 有很多 JIT。Pyston 和 Pyjion 就是 JIT。

为 Python 3.13 提议的 JIT 是一个任意复制补丁（copy-and-patch）的 JIT。

什么是复制和补丁（copy-and-patch） JIT？
没听说过Copy-and-patch JIT？别担心，我和大多数人都没听说过。它是最近才在 2021 年提出的一个想法，旨在作为动态语言运行时的快速算法。

我将尝试通过扩展我们的解释器循环并将其重写为 JIT 来解释什么是复制和补丁（copy-and-patch） JIT。之前，解释器循环做了两件事，首先是解释（查看字节码），然后是执行（运行指令）。我们可以做的是将这些任务分开，让解释器输出指令，而不是执行指令。

复制和补丁（copy-and-patch）JIT 的原理是，复制每条命令的指令，并为字节码参数（或补丁）填空。下面是一个重写的示例，我保持了非常相似的循环，但每次都附加了一个要执行的 Python 代码字符串：

def copy_and_patch_interpret(bytecodes):
    code = 'def f():\n'
    code += '  stack = []\n'
    code += '  variables = {}\n'
    for bytecode, arg in bytecodes:
        if bytecode == "LOAD_CONST":
            code += f'  stack.append({arg})\n'
        elif bytecode == "LOAD_FAST":
            code += f'  stack.append(variables["{arg}"])\n'
        elif bytecode == "STORE_FAST":
            code += f'  variables["{arg}"] = stack.pop()\n'
        elif bytecode == "RETURN_VALUE":
            code += '  return stack.pop()\n'
    code += 'f()'
    return code

原始函数的结果是

def f():
  stack = []
  variables = {}
  stack.append(1)
  variables["a"] = stack.pop()
  stack.append(None)
  return stack.pop()
f()


这一次，代码是连续的，不需要循环执行。我们可以存储生成的字符串，然后运行任意多次：

compiled_function = compile(copy_and_patch_interpret(func), filename="<string>", mode="exec")

print(exec(compiled_function))
print(exec(compiled_function))
print(exec(compiled_function))

这有什么意义呢？结果代码做的是同样的事情，但运行速度应该更快。

为什么要使用复制和修补 JIT？
与 "完整 "的 JIT 编译器相比，这种为每个字节码编写指令并修补值的技术各有利弊。

完整的 JIT 编译器通常会将 LOAD_FAST 这样的高级字节码编译成 IL（中间语言）中的低级指令。由于每种 CPU 架构都有不同的指令和功能，要编写一个能将高级代码直接转换为机器代码的编译器，并支持 32 位和 64 位 CPU，以及苹果的 ARM 架构和所有其他类型的 ARM，是一件非常复杂的事情。

相反，大多数 JIT 首先编译的是 IL，即类似于通用机器码的指令集。这些指令包括 "PUSH 一个 64 位整数"、"POP 一个 64 位浮点数"、"MULTIPLY 堆栈中的值 "等。

然后，JIT 可以在运行时将 IL 编译成机器码，方法是发出特定于 CPU 的指令，并将其存储在内存中以便以后执行（类似于我们在示例中的做法）。

一旦有了 IL，你就可以对代码进行各种有趣的优化，比如常量传播和循环提升。你可以在 Pyjion 的实时编译器 UI 中看到一个例子。

完全 "JIT 的最大缺点是：编译成 IL 后再编译成机器代码的过程非常缓慢。不仅速度慢，而且占用大量内存。

为了说明这一点，我们可以从最近的研究 "Python 与 JIT 编译器的结合：一个简单的实现和一个比较评估"中的数据显示，基于 Java 的 Python JIT（如 GraalPy 和 Jython）的启动时间比普通 CPython 长 100 倍，编译时需要消耗额外的 Gigabyte 内存。目前已经有针对 Python 的完整 JIT 实现。

之所以选择 "复制加修补Copy-and-patch"，是因为字节码到机器码的编译是以一组 "模板 "的形式完成的，这些模板在运行时被拼接在一起，并用正确的值进行修补。

这意味着普通 Python 用户不会在 Python 运行时中运行这种复杂的 JIT 编译器架构。

Python 编写自己的 IL 和 JIT 也是不合理的，因为像 LLVM 和 ryuJIT 这样的现成编译器已经很多了。

但完整的 JIT 需要将这些工具与 Python 捆绑在一起，并增加所有开销。

Copy-and-patch JIT 只需要在编译 CPython 源代码的机器上安装 LLVM JIT 工具，对于大多数人来说，这意味着为 python.org 编译和打包 CPython 的 CI 机器。

那么这个 JIT 是如何工作的呢？
Python 的复制和补丁编译器是通过在 Python 3.13 的 API 中扩展一些新的（老实说并不广为人知的）API 来工作的。这些变化使得 CPython 在运行时可以发现可插拔的优化器，并控制代码的执行方式。这个新的 JIT 是这个新架构的可选优化器。我认为，一旦主要的 bug 被解决，它将成为未来版本的默认优化器。

从源代码编译 CPython 时，可以在 configure 脚本中提供一个标志--enable-experimental-jit。这将为 Python 字节码生成机器码模板。首先复制每个字节码的 C 代码，例如最简单的 LOAD_CONST：

frame->instr_ptr = next_instr;
next_instr += 1;
INSTRUCTION_STATS(LOAD_CONST); // Not used unless compiled with instrumentation
PyObject *value;
value = GETITEM(FRAME_CO_CONSTS, oparg);
Py_INCREF(value);
stack_pointer[0] = value;
stack_pointer += 1;
DISPATCH();

这种字节码的指令首先由 C 语言编译成一个小的共享库，然后存储为机器码。由于有些变量（如 oparg）通常在运行时确定，因此 C 代码在编译时会将这些参数留为 0。就 LOAD_CONST 而言，有 2 个孔需要填入，即 oparg 和下一条指令：

static const Hole _LOAD_CONST_code_holes[3] = {
    {0xd, HoleKind_X86_64_RELOC_UNSIGNED, HoleValue_OPARG, NULL, 0x0},
    {0x46, HoleKind_X86_64_RELOC_UNSIGNED, HoleValue_CONTINUE, NULL, 0x0},
};

然后，所有机器码都会以字节序列的形式保存在 jit_stencil.h 文件中，该文件会在新的编译阶段自动生成。反汇编代码以注释的形式保存在每个字节码模板的上方，其中 JIT_OPARG 和 JIT_CONTINUE 是需要填补的漏洞：

0000000000000000 <__JIT_ENTRY>:
pushq   %rbp
movq    %rsp, %rbp
movq    (%rdi), %rax
movq    0x28(%rax), %rax
movabsq $0x0, %rcx
000000000000000d:  X86_64_RELOC_UNSIGNED        __JIT_OPARG
movzwl  %cx, %ecx
movq    0x28(%rax,%rcx,8), %rax
movl    0xc(%rax), %ecx
incl    %ecx
je      0x3d <__JIT_ENTRY+0x3d>
movq    %gs:0x0, %r8
cmpq    (%rax), %r8
jne     0x37 <__JIT_ENTRY+0x37>
movl    %ecx, 0xc(%rax)
jmp     0x3d <__JIT_ENTRY+0x3d>
lock
addq    $0x4, 0x10(%rax)
movq    %rax, (%rsi)
addq    $0x8, %rsi
movabsq $0x0, %rax
0000000000000046:  X86_64_RELOC_UNSIGNED        __JIT_CONTINUE
popq    %rbp
jmpq    *%rax

新的 JIT 编译器启动后，会将每个字节码的机器码指令复制到一个序列中，并将每个模板的值替换为代码对象中该字节码的参数。生成的机器码存储在内存中，然后每次运行 Python 函数时，都会直接执行该机器码。

如果您编译我的分支并在测试脚本上试用，然后将其交给 Ada Pro 或 Hopper 等反汇编器，就能看到 JIT 化的代码。目前，只有在函数包含 JUMP_BACKWARD 操作码（用于 while 语句）时才会使用 JIT，但将来会有所改变。

更快吗？
最初的基准测试显示性能提高了 5-9%。

提前编译现有解释器的挑战在于，进行认真优化的机会较少。Python 3.11 的自适应解释器是朝着正确方向迈出的一步，但 Python 还需要走得更远才能看到性能的阶跃变化。