关于生成器的说明

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.3yield $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 点,我们已经计算了两组变量:

  • live variables 是指在恢复点之后的代码(即在 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。其他值表示从 1 开始的 yield 点的索引(对应于上面生成器信息中显示的索引)。

  • 第二个成员,即 参数结构 在首次初始化后是只读的。它存储了生成器函数被调用时的参数值。在我们的例子中,这些是 xy 的值。

  • 第三个成员,即 状态结构 ,存储如上计算的实时变量。

具体来说,我们的示例生成器结构(假设生成器函数使用浮点数调用)如下:

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;
   }
}

注意,这里保存 xy 是多余的:Numba 无法识别状态变量 xyarg0arg1 具有相同的值。

分配

Numba 如何确保生成器结构被保留足够长的时间?有两种情况:

  • 当从Numba编译的函数调用Numba编译的生成器函数时,结构由被调用者分配在栈上。在这种情况下,生成器的实例化实际上是无成本的。

  • 当从常规Python代码调用Numba编译的生成器函数时,会实例化一个与CPython兼容的包装器,该包装器具有存储结构所需的适当空间,并且其 tp_iternext 槽是生成器原生代码的包装器。

编译为本机代码

在编译生成器函数时,Numba 实际上生成了三个原生函数:

  • 一个初始化函数。这是与生成器函数本身对应的函数:它接收函数参数并将它们存储在生成器结构中(通过指针传递)。它还将 恢复索引 初始化为0,表示生成器尚未启动。

  • 一个 next() 函数。这是在生成器内部恢复执行时调用的函数。它的单个参数是指向生成器结构的指针,并返回下一个产生的值(如果生成器耗尽,则返回一个特殊的退出代码,以便在从 Numba 编译的函数调用时快速检查)。

  • 一个可选的终结器。在对象模式下,此函数确保生成器状态中存储的所有活动变量都被解引用,即使生成器在未被完全使用的情况下被销毁。

next() 函数

next() 函数是这三个原生函数中最不直接的一个。它从一个蹦床开始,根据存储在生成器结构中的 恢复索引 将执行分派到正确的恢复点。以下是该函数在我们的示例中可能的开始方式:

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。跳板开关有三个目标(每个 resume index 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 表示产生了一个值,返回 -3 表示 StopIteration。%rcx 指向生成器结构的开始,其中存储了恢复索引。