7.4。关于发电机的注意事项

7.4。关于发电机的注意事项

原文： http://numba.pydata.org/numba-doc/latest/developer/generators.html

Numba 最近获得了编译发电机功能的支持。本文档解释了一些实现选择。

7.4.1。术语

为清楚起见，我们区分 发生器功能 和 发生器 。生成器函数是包含一个或多个yield语句的函数。生成器（有时也称为“生成器迭代器”）是生成器函数的返回值;每次调用 next() 时，它会在其帧内恢复执行。

屈服点 是调用yield语句的地方。 恢复点 就在 屈服点 之后的位置，其中当再次调用 next() 时恢复执行。

7.4.2。功能分析

假设我们有以下简单的生成器函数：

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 $0.3和yield $0.7）。

第二部分显示了特定于发电机的信息。要理解它，我们必须了解暂停和恢复生成器的含义。

挂起生成器时，我们不仅仅向调用者返回一个值（yield语句的操作数）。我们还必须保存发生器的 当前状态 以便恢复执行。在简单的用例中，可能会保留 CPU 的寄存器值或堆栈槽，直到下一次调用 next（）。但是，任何非平凡的案例都会无可救药地破坏这些价值观，因此我们必须将它们保存在一个定义明确的地方。

我们需要保存哪些值？那么，在 Numba Intermediate Representation 的背景下，我们必须在每个屈服点保存所有 实时变量 。由于控制流程图，计算了这些实时变量。

保存实时变量并暂停生成器后，恢复生成器只需执行逆操作：实时变量将从保存的生成器状态恢复。

注意

这是相同的分析，有助于在适当的地方插入 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语句之后）
弱活变量 是在恢复点之后立即进行定义的变量;它们必须保存在对象模式中，以确保正确的参考清理

7.4.3。发电机结构

7.4.3.1。布局

功能分析有助于我们收集足够的信息来定义生成器结构的布局，该布局将存储生成器的整个执行状态。这是生成器结构布局的草图，用伪代码表示：

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。其他值表示屈服点的指数从 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具有相同的值。

7.4.3.2。分配

Numba 如何确保发电机结构保持足够长的时间？有两种情况：

当从 Numba 编译的函数调用 Numba 编译的生成器函数时，该结构由被调用者在堆栈上分配。在这种情况下，发电机实例化实际上是无成本的。
当从常规 Python 代码调用 Numba 编译的生成器函数时，实例化 CPython 兼容的包装器，其具有适当的分配空间来存储结构，并且其 tp_iternext 插槽是一个包装器生成器的本机代码。

7.4.4。编译为本机代码

在编译生成器函数时，Numba 实际生成了三个本机函数：

初始化函数。这是与生成器函数本身对应的函数：它接收函数参数并将它们存储在生成器结构（由指针传递）中。它还将 恢复指数 初始化为 0，表明发电机尚未启动。
next（）函数。这是在生成器内恢复执行的函数。它的单个参数是指向生成器结构的指针，它返回下一个产生的值（如果生成器耗尽，则使用特殊的退出代码，以便在从 Numba 编译的函数调用时进行快速检查）。
可选的终结器。在对象模式下，此功能可确保存储在生成器状态中的所有实时变量都被减少，即使生成器在没有耗尽的情况下被销毁也是如此。

7.4.4.1。 next（）函数

next（）函数是三个本机函数中最不直接的函数。它以蹦床开始，根据存储在发生器结构中的 恢复索引 ，将执行分配到右恢复点。以下是函数 start 在我们的示例中的外观：

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）是第一个（相应的第二个）屈服点之后的恢复点。

在 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 表示产生一个值，-3 表示 StopIteration。 %rcx指向生成恢复索引的生成器结构的开始。