生成器说明
Numba 最近增加了对编译生成器函数的支持。本文档解释了一些实现选择。
术语
为清晰起见，我们区分 *生成器函数* 和 *生成器*。生成器函数是包含一个或多个 yield 语句的函数。生成器（有时也称为“生成器迭代器”）是生成器函数的返回值；每次调用 next() 时，它都会在其帧内恢复执行。
*yield点* 是调用 yield 语句的地方。*恢复点* 是紧接在 *yield点* 之后的地方，当再次调用 next() 时，执行在此处恢复。
函数分析
假设我们有以下简单的生成器函数
def gen(x, y):
yield x + y
yield x - y
这是它的 CPython 字节码，使用 dis.dis() 打印输出
7 0 LOAD_FAST 0 (x)
3 LOAD_FAST 1 (y)
6 BINARY_ADD
7 YIELD_VALUE
8 POP_TOP
8 9 LOAD_FAST 0 (x)
12 LOAD_FAST 1 (y)
15 BINARY_SUBTRACT
16 YIELD_VALUE
17 POP_TOP
18 LOAD_CONST 0 (None)
21 RETURN_VALUE
当使用 NUMBA_DUMP_IR 设置为 1 编译此函数时，会打印出以下信息
----------------------------------IR DUMP: gen----------------------------------
label 0:
x = arg(0, name=x) ['x']
y = arg(1, name=y) ['y']
$0.3 = x + y ['$0.3', 'x', 'y']
$0.4 = yield $0.3 ['$0.3', '$0.4']
del $0.4 []
del $0.3 []
$0.7 = x - y ['$0.7', 'x', 'y']
del y []
del x []
$0.8 = yield $0.7 ['$0.7', '$0.8']
del $0.8 []
del $0.7 []
$const0.9 = const(NoneType, None) ['$const0.9']
$0.10 = cast(value=$const0.9) ['$0.10', '$const0.9']
del $const0.9 []
return $0.10 ['$0.10']
------------------------------GENERATOR INFO: gen-------------------------------
generator state variables: ['$0.3', '$0.7', 'x', 'y']
yield point #1: live variables = ['x', 'y'], weak live variables = ['$0.3']
yield point #2: live variables = [], weak live variables = ['$0.7']
这是什么意思？第一部分是 Numba IR，如 阶段 2：生成 Numba IR 中所示。我们可以看到两个 yield点（yield $0.3 和 yield $0.7）。
第二部分显示了生成器特有的信息。要理解它，我们必须理解暂停和恢复生成器意味着什么。
暂停生成器时，我们不仅仅是向调用者返回一个值（yield 语句的操作数）。我们还需要保存生成器的 *当前状态*，以便恢复执行。在简单用例中，CPU 的寄存器值或栈槽或许可以保留到下一次调用 next()。然而，任何非简单用例都会不可避免地破坏这些值，因此我们必须将它们保存在一个明确定义的位置。
我们需要保存哪些值？在 Numba 中间表示的上下文中，我们必须在每个 yield点保存所有 *活跃变量*。这些活跃变量通过控制流图计算得出。
一旦活跃变量被保存并且生成器暂停，恢复生成器就只是一个逆向操作：活跃变量从保存的生成器状态中恢复。
注意
同样的分析也有助于在适当的地方插入 Numba del 指令。
让我们再次回顾生成器信息
generator state variables: ['$0.3', '$0.7', 'x', 'y']
yield point #1: live variables = ['x', 'y'], weak live variables = ['$0.3']
yield point #2: live variables = [], weak live variables = ['$0.7']
Numba 已经计算出所有活跃变量的并集（表示为“状态变量”）。这将有助于定义 生成器结构 的布局。此外，对于每个 yield点，我们计算了两组变量
*活跃变量* 是恢复点之后（即
yield语句之后）代码所使用的变量*弱活跃变量* 是在恢复点之后立即被删除的变量；它们必须在 对象模式 下保存，以确保正确的引用清理。
生成器结构
布局
函数分析帮助我们收集足够的信息来定义生成器结构的布局，该结构将存储生成器的整个执行状态。这是生成器结构布局的伪代码草图
struct gen_struct_t {
int32_t resume_index;
struct gen_args_t {
arg_0_t arg0;
arg_1_t arg1;
...
arg_N_t argN;
}
struct gen_state_t {
state_0_t state_var0;
state_1_t state_var1;
...
state_N_t state_varN;
}
}
让我们按顺序描述这些字段。
第一个成员，即 *恢复索引*，是一个整数，指示生成器必须在哪个恢复点恢复执行。按照惯例，它可能有两个特殊值：0 表示执行必须从生成器开头开始（即第一次调用
next()时）；-1 表示生成器已耗尽，恢复必须立即引发 StopIteration。其他值表示 yield点 的索引，从 1 开始（对应于上面生成器信息中显示的索引）。第二个成员，即 *参数结构*，在首次初始化后是只读的。它存储了调用生成器函数时所用的参数值。在我们的示例中，这些是
x和y的值。第三个成员，即 *状态结构*，存储了如上计算的活跃变量。
具体而言，我们示例中的生成器结构（假设生成器函数使用浮点数调用）如下
struct gen_struct_t {
int32_t resume_index;
struct gen_args_t {
double arg0;
double arg1;
}
struct gen_state_t {
double $0.3;
double $0.7;
double x;
double y;
}
}
请注意，在此处保存 x 和 y 是多余的：Numba 无法识别状态变量 x 和 y 与 arg0 和 arg1 具有相同的值。
分配
Numba 如何确保生成器结构被长时间保留？有两种情况
当从 Numba 编译的函数中调用 Numba 编译的生成器函数时，结构由被调用者在栈上分配。在这种情况下，生成器实例化几乎没有成本。
当从常规 Python 代码调用 Numba 编译的生成器函数时，会实例化一个 CPython 兼容的包装器，该包装器具有足够的分配空间来存储结构，并且其
tp_iternext槽是对生成器原生代码的包装。
编译为原生代码
编译生成器函数时，Numba 实际上会生成三个原生函数
一个初始化函数。这是对应于生成器函数本身的函数：它接收函数参数并将它们存储在生成器结构中（通过指针传递）。它还将 *恢复索引* 初始化为 0，表示生成器尚未开始。
一个 next() 函数。这是在生成器内部恢复执行时调用的函数。它的唯一参数是指向生成器结构的指针，它返回下一个 yield 的值（如果生成器已耗尽，则使用特殊退出代码，以便从 Numba 编译的函数调用时快速检查）。
一个可选的终结器。在对象模式下，此函数确保存储在生成器状态中的所有活跃变量都被 decref（引用计数递减），即使生成器在未耗尽的情况下被销毁。
next() 函数
next() 函数是这三个原生函数中最不直接的。它以一个跳转器（trampoline）开始，根据存储在生成器结构中的 *恢复索引* 将执行分派到正确的恢复点。以下是函数在我们示例中可能的样子
define i32 @"__main__.gen.next"(
double* nocapture %retptr,
{ i8*, i32 }** nocapture readnone %excinfo,
i8* nocapture readnone %env,
{ i32, { double, double }, { double, double, double, double } }* nocapture %arg.gen)
{
entry:
%gen.resume_index = getelementptr { i32, { double, double }, { double, double, double, double } }* %arg.gen, i64 0, i32 0
%.47 = load i32* %gen.resume_index, align 4
switch i32 %.47, label %stop_iteration [
i32 0, label %B0
i32 1, label %generator_resume1
i32 2, label %generator_resume2
]
; rest of the function snipped
（为提高可读性，从 LLVM IR 中删减了不相关内容）
我们在 %arg.gen 中识别出指向生成器结构的指针。跳转器开关有三个目标（每个 *恢复索引* 0、1 和 2 各一个），以及一个名为 stop_iteration 的备用目标标签。标签 B0 表示函数开始，generator_resume1（或 generator_resume2）是第一个（或第二个）yield点之后的恢复点。
经 LLVM 生成后，该函数的完整原生汇编代码可能如下所示（在 x86-64 上）
.globl __main__.gen.next
.align 16, 0x90
__main__.gen.next:
movl (%rcx), %eax
cmpl $2, %eax
je .LBB1_5
cmpl $1, %eax
jne .LBB1_2
movsd 40(%rcx), %xmm0
subsd 48(%rcx), %xmm0
movl $2, (%rcx)
movsd %xmm0, (%rdi)
xorl %eax, %eax
retq
.LBB1_5:
movl $-1, (%rcx)
jmp .LBB1_6
.LBB1_2:
testl %eax, %eax
jne .LBB1_6
movsd 8(%rcx), %xmm0
movsd 16(%rcx), %xmm1
movaps %xmm0, %xmm2
addsd %xmm1, %xmm2
movsd %xmm1, 48(%rcx)
movsd %xmm0, 40(%rcx)
movl $1, (%rcx)
movsd %xmm2, (%rdi)
xorl %eax, %eax
retq
.LBB1_6:
movl $-3, %eax
retq
请注意，函数返回 0 表示 yield 了一个值，返回 -3 表示 StopIteration。%rcx 指向生成器结构的起始处，恢复索引存储在那里。