Numba 架构

介绍

Numba 是一个带有可选类型特化的 Python 字节码编译器。

假设你将一个函数输入到标准的 Python 解释器中(以下简称“CPython”):

def add(a, b):
    return a + b

解释器会立即解析该函数,并将其转换为字节码表示形式,该表示形式描述了CPython解释器应在低级别上如何执行该函数。对于上面的示例,它看起来像这样:

>>> import dis
>>> dis.dis(add)
2           0 LOAD_FAST                0 (a)
            3 LOAD_FAST                1 (b)
            6 BINARY_ADD
            7 RETURN_VALUE

CPython 使用基于堆栈的解释器(很像 HP 计算器),因此代码首先将两个局部变量压入堆栈。BINARY_ADD 操作码从堆栈中弹出前两个参数,并调用一个等同于调用 a.__add__(b) 的 Python C API 函数。结果随后被压入解释器堆栈的顶部。最后,RETURN_VALUE 操作码将堆栈顶部的值作为函数调用的结果返回。

Numba 可以获取这个字节码并将其编译为机器码,该机器码执行与 CPython 解释器相同的操作,将 ab 视为通用 Python 对象。Python 的完整语义被保留,并且编译后的函数可以与任何定义了加法运算符的对象一起使用。当 Numba 函数以这种方式编译时,我们称其已在 对象模式 下编译,因为代码仍然操作 Python 对象。

在对象模式下编译的Numba代码并不比在CPython解释器中执行原始Python函数快多少。然而,如果我们专门化函数以仅使用某些数据类型运行,Numba可以生成更短且更高效的代码,这些代码可以原生地操作数据,而无需调用Python C API。当代码已针对特定数据类型进行编译,使得函数体不再依赖于Python运行时,我们称该函数已在 nopython模式 下编译。在nopython模式下编译的数值代码可以比原始Python代码快数百倍。

编译器架构

与许多编译器类似,Numba 在概念上可以分为 前端后端

Numba 的 前端 包括分析 Python 字节码、将其翻译为 Numba IR 并在 IR 上执行各种转换和分析步骤的阶段。关键步骤之一是 类型推断 。为了使后端能够在 nopython 模式 下生成代码,前端必须成功地对所有变量进行明确的类型推断,因为后端使用类型信息来匹配适当的代码生成器与它们操作的值。

Numba 的 后端 遍历由前端分析生成的 Numba IR,并利用类型推断阶段推导出的类型信息,为每个遇到的运算生成正确的 LLVM 代码。在生成 LLVM 代码后,会请求 LLVM 库对其进行优化,并为最终的本地函数生成本地处理器代码。

除了编译器的前端和后端之外,还有其他部分,例如用于JIT函数的缓存机制。本文档不考虑这些部分。

上下文

Numba 非常灵活,允许它为不同的硬件架构(如CPU和GPU)生成代码。为了支持这些不同的应用,Numba 使用了一个 类型上下文 和一个 目标上下文

在编译器前端中,类型上下文 用于对函数中的操作和值执行类型推断。类似的类型上下文可以用于许多架构,因为在几乎所有情况下,类型推断都是与硬件无关的。然而,Numba 目前为每个目标使用不同的类型上下文。

一个 目标上下文 用于生成在类型推断期间识别的Numba类型所需的具体指令序列。目标上下文是特定于架构的,并且在定义执行模型和可用的Python API方面具有灵活性。例如,Numba有一个针对这两种架构的“cpu”和“cuda”上下文,以及一个生成多线程CPU代码的“parallel”上下文。

编译器阶段

Numba 中的 jit() 装饰器最终调用 numba.compiler.compile_extra() ,该函数通过以下描述的多阶段过程编译 Python 函数。

阶段 1: 分析字节码

在编译开始时,函数字节码被传递给 Numba 解释器 (numba.interpreter) 的一个实例。解释器对象分析字节码以找到控制流图 (numba.controlflow)。控制流图 (CFG) 描述了由于循环和分支,执行如何在函数内的一个块移动到下一个块的方式。

数据流分析 (numba.dataflow) 接受控制流图并追踪不同代码路径下值如何被推入和弹出Python解释器栈。这对于理解栈上变量的生命周期至关重要,这是在第二阶段所需要的。

如果你将环境变量 NUMBA_DUMP_CFG 设置为 1,Numba 会将控制流图分析的结果输出到屏幕上。我们的 add() 示例相当简单,因为只有一个语句块:

CFG adjacency lists:
{0: []}
CFG dominators:
{0: set([0])}
CFG post-dominators:
{0: set([0])}
CFG back edges: []
CFG loops:
{}
CFG node-to-loops:
{0: []}

一个具有更复杂流程控制的函数将有一个更有趣的控制流图。这个函数:

def doloops(n):
    acc = 0
    for i in range(n):
        acc += 1
        if n == 10:
            break
    return acc

编译为此字节码:

 9           0 LOAD_CONST               1 (0)
             3 STORE_FAST               1 (acc)

10           6 SETUP_LOOP              46 (to 55)
             9 LOAD_GLOBAL              0 (range)
            12 LOAD_FAST                0 (n)
            15 CALL_FUNCTION            1
            18 GET_ITER
       >>   19 FOR_ITER                32 (to 54)
            22 STORE_FAST               2 (i)

11          25 LOAD_FAST                1 (acc)
            28 LOAD_CONST               2 (1)
            31 INPLACE_ADD
            32 STORE_FAST               1 (acc)

12          35 LOAD_FAST                0 (n)
            38 LOAD_CONST               3 (10)
            41 COMPARE_OP               2 (==)
            44 POP_JUMP_IF_FALSE       19

13          47 BREAK_LOOP
            48 JUMP_ABSOLUTE           19
            51 JUMP_ABSOLUTE           19
       >>   54 POP_BLOCK

14     >>   55 LOAD_FAST                1 (acc)
            58 RETURN_VALUE

此字节码对应的 CFG 是:

CFG adjacency lists:
{0: [6], 6: [19], 19: [54, 22], 22: [19, 47], 47: [55], 54: [55], 55: []}
CFG dominators:
{0: set([0]),
 6: set([0, 6]),
 19: set([0, 6, 19]),
 22: set([0, 6, 19, 22]),
 47: set([0, 6, 19, 22, 47]),
 54: set([0, 6, 19, 54]),
 55: set([0, 6, 19, 55])}
CFG post-dominators:
{0: set([0, 6, 19, 55]),
 6: set([6, 19, 55]),
 19: set([19, 55]),
 22: set([22, 55]),
 47: set([47, 55]),
 54: set([54, 55]),
 55: set([55])}
CFG back edges: [(22, 19)]
CFG loops:
{19: Loop(entries=set([6]), exits=set([54, 47]), header=19, body=set([19, 22]))}
CFG node-to-loops:
{0: [], 6: [], 19: [19], 22: [19], 47: [], 54: [], 55: []}

CFG 中的数字指的是上述操作码名称左侧显示的字节码偏移量。

阶段 2:生成 Numba IR

一旦控制流和数据分析完成,Numba解释器可以逐步执行字节码并将其翻译成Numba内部的中间表示。这个翻译过程将函数从栈机表示(由Python解释器使用)转换为寄存器机表示(由LLVM使用)。

尽管IR作为对象树存储在内存中,但它可以序列化为字符串以供调试。如果你将环境变量 NUMBA_DUMP_IR 设置为1,Numba IR将被转储到屏幕上。对于上述描述的 add() 函数,Numba IR看起来像:

label 0:
    a = arg(0, name=a)                       ['a']
    b = arg(1, name=b)                       ['b']
    $0.3 = a + b                             ['$0.3', 'a', 'b']
    del b                                    []
    del a                                    []
    $0.4 = cast(value=$0.3)                  ['$0.3', '$0.4']
    del $0.3                                 []
    return $0.4                              ['$0.4']

del 指令由 实时变量分析 生成。这些指令确保引用不会泄漏。在 nopython 模式 中,一些对象由 Numba 运行时跟踪,而另一些则不跟踪。对于被跟踪的对象,会发出一个解引用操作;否则,该指令是一个空操作。在 对象模式 中,每个变量包含一个对 PyObject 的所有引用。

阶段 3:重写无类型 IR

在运行类型推断之前,可能需要在Numba IR上运行某些转换。一个这样的例子是检测具有隐式常量参数的 raise 语句,以便在 nopython模式 中支持它们。假设你用Numba编译以下函数:

def f(x):
   if x == 0:
      raise ValueError("x cannot be zero")

如果你将 NUMBA_DUMP_IR 环境变量设置为 1,你将看到在类型推断阶段之前被重写的IR:

REWRITING:
    del $0.3                                 []
    $12.1 = global(ValueError: <class 'ValueError'>) ['$12.1']
    $const12.2 = const(str, x cannot be zero) ['$const12.2']
    $12.3 = call $12.1($const12.2)           ['$12.1', '$12.3', '$const12.2']
    del $const12.2                           []
    del $12.1                                []
    raise $12.3                              ['$12.3']
____________________________________________________________
    del $0.3                                 []
    $12.1 = global(ValueError: <class 'ValueError'>) ['$12.1']
    $const12.2 = const(str, x cannot be zero) ['$const12.2']
    $12.3 = call $12.1($const12.2)           ['$12.1', '$12.3', '$const12.2']
    del $const12.2                           []
    del $12.1                                []
    raise <class 'ValueError'>('x cannot be zero') []

阶段 4:推断类型

既然已经生成了 Numba IR,就可以进行类型分析了。函数参数的类型可以从 @jit 装饰器中给出的显式函数签名中获取(例如 @jit('float64(float64, float64)')),或者如果编译发生在函数首次调用时,则可以从实际函数参数的类型中获取。

类型推断引擎位于 numba.typeinfer 中。它的任务是为 Numba IR 中的每个中间变量分配一个类型。可以通过将 NUMBA_DUMP_ANNOTATION 环境变量设置为 1 来查看此过程的结果:

-----------------------------------ANNOTATION-----------------------------------
# File: archex.py
# --- LINE 4 ---

@jit(nopython=True)

# --- LINE 5 ---

def add(a, b):

    # --- LINE 6 ---
    # label 0
    #   a = arg(0, name=a)  :: int64
    #   b = arg(1, name=b)  :: int64
    #   $0.3 = a + b  :: int64
    #   del b
    #   del a
    #   $0.4 = cast(value=$0.3)  :: int64
    #   del $0.3
    #   return $0.4

    return a + b

如果类型推断无法为所有中间变量找到一致的类型分配,它会将每个变量标记为类型 pyobject 并回退到对象模式。当在函数体中使用不支持的Python类型、语言特性或函数时,类型推断可能会失败。

备注

截至 Numba 0.59,对象模式回退仅在启用 loop-lifting 时发生。

阶段 5a:重写类型化 IR

此过程的目的是执行任何仍然需要,或至少可以从 Numba IR 类型信息中受益的高级优化。

一个不容易在降低级别后优化的领域示例是多维数组操作领域。当Numba降低一个数组操作时,Numba将该操作视为一个完整的ufunc内核。在降低单个数组操作期间,Numba生成一个内联广播循环,该循环创建一个新的结果数组。然后Numba生成一个应用循环,该循环在数组输入上应用该操作符。一旦这些循环被降低到LLVM中,识别和重写它们是困难的,如果不是不可能的话。

在数组操作符领域的一个优化例子是循环融合和快捷消枝。当优化器识别到一个数组操作符的输出被直接输入到另一个数组操作符,并且仅输入到该操作符时,它可以将这两个循环融合为一个循环。优化器可以通过直接将第一个操作的结果输入到第二个操作中,跳过对中间数组的存储和加载,从而进一步消除为初始操作分配的临时数组。这种消除被称为快捷消枝。Numba 目前使用重写传递来实现这些数组优化。更多信息,请参阅本文档后面的“案例研究:数组表达式”小节。

可以通过将 NUMBA_DUMP_IR 环境变量设置为非零值(例如1)来查看重写的结果。以下示例显示了重写过程的输出,因为它识别出一个由乘法和加法组成的数组表达式,并输出一个融合内核作为特殊操作符, arrayexpr():

______________________________________________________________________
REWRITING:
a0 = arg(0, name=a0)                     ['a0']
a1 = arg(1, name=a1)                     ['a1']
a2 = arg(2, name=a2)                     ['a2']
$0.3 = a0 * a1                           ['$0.3', 'a0', 'a1']
del a1                                   []
del a0                                   []
$0.5 = $0.3 + a2                         ['$0.3', '$0.5', 'a2']
del a2                                   []
del $0.3                                 []
$0.6 = cast(value=$0.5)                  ['$0.5', '$0.6']
del $0.5                                 []
return $0.6                              ['$0.6']
____________________________________________________________
a0 = arg(0, name=a0)                     ['a0']
a1 = arg(1, name=a1)                     ['a1']
a2 = arg(2, name=a2)                     ['a2']
$0.5 = arrayexpr(ty=array(float64, 1d, C), expr=('+', [('*', [Var(a0, test.py (14)), Var(a1, test.py (14))]), Var(a2, test.py (14))])) ['$0.5', 'a0', 'a1', 'a2']
del a0                                   []
del a1                                   []
del a2                                   []
$0.6 = cast(value=$0.5)                  ['$0.5', '$0.6']
del $0.5                                 []
return $0.6                              ['$0.6']
______________________________________________________________________

在此重写之后,Numba 将数组表达式降低为一个类似 ufunc 的新函数,该函数被内联到一个仅分配单个结果数组的单个循环中。

阶段 5b:执行自动并行化

只有在 jit() 装饰器中的 parallel 选项设置为 True 时,才会执行此过程。此过程在 Numba IR 的操作语义中寻找隐含的并行性,并使用特殊的 parfor 操作符将这些操作替换为显式的并行表示。然后,执行优化以最大化相邻 parfor 的数量,以便它们可以融合在一起,形成一个只需一次数据遍历的 parfor,从而通常具有更好的缓存性能。最后,在降低级别时,这些 parfor 操作符被转换为类似于 guvectorize 的形式,以实现实际的并行性。

自动并行化过程包含多个子过程,其中许多可以通过传递给 jit()parallel 关键字参数的选项字典来控制:

{ 'comprehension': True/False,  # parallel comprehension
  'prange':        True/False,  # parallel for-loop
  'numpy':         True/False,  # parallel numpy calls
  'reduction':     True/False,  # parallel reduce calls
  'setitem':       True/False,  # parallel setitem
  'stencil':       True/False,  # parallel stencils
  'fusion':        True/False,  # enable fusion or not
}

它们的默认值都设置为 True。子过程在以下段落中有更详细的描述。

  1. CFG 简化

    有时,Numba IR 会包含一系列没有循环的块,这些块在此子过程中被合并为单个块。此子过程简化了后续的 IR 分析。

  2. Numpy 规范化

    一些 Numpy 操作可以写成对 Numpy 对象的操作(例如 arr.sum()),或者作为对这些对象的 Numpy 调用(例如 numpy.sum(arr))。此子过程将所有此类操作转换为后一种形式,以便进行更清晰的后续分析。

  3. 数组分析

    后续 parfor 融合的一个关键要求是 parfor 具有相同的迭代空间,这些迭代空间通常对应于 Numpy 数组的维度大小。在此子过程中,分析 IR 以确定 Numpy 数组维度的等价类。考虑示例 a = b + 1,其中 ab 都是 Numpy 数组。在这里,我们知道 a 的每个维度必须与 b 的相应维度具有相同的等价类。通常,富含 Numpy 操作的例程将使所有在函数内创建的数组的等价类完全已知。

    数组分析还将推断切片选择的大小等价性,以及布尔数组掩码(仅限一维)。例如,它能够推断出 a[1 : n-1] 的大小与 b[0 : n-2] 相同。

    数组分析也可能插入安全假设,以确保在操作可以并行化之前满足与数组大小相关的先决条件。例如,np.dot(X, w) 在二维矩阵 X 和一维向量 w 之间,要求 X 的第二个维度与 w 的大小相同。通常这种运行时检查会自动插入,但如果数组分析能够推断出这种等价性,它将跳过这些检查。

    用户甚至可以通过将关于数组大小的隐性知识转化为显性断言来帮助数组分析。例如,在下面的代码中:

    @numba.njit(parallel=True)
    def logistic_regression(Y, X, w, iterations):
        assert(X.shape == (Y.shape[0], w.shape[0]))
        for i in range(iterations):
            w -= np.dot(((1.0 / (1.0 + np.exp(-Y * np.dot(X, w))) - 1.0) * Y), X)
        return w
    

    明确地进行断言有助于消除函数其余部分中的所有边界检查。

  4. prange() 到 parfor

    在 for 循环中使用 prange (显式并行循环) 是程序员明确指示该 for 循环的所有迭代可以并行执行。在此子过程中,我们分析控制流图 (CFG) 以定位循环,并将那些由 prange 对象控制的循环转换为显式的 parfor 操作符。每个显式的 parfor 操作符包括:

    1. 一个描述 parfor 迭代空间的循环嵌套信息列表。循环嵌套列表中的每个条目包含一个索引变量、范围的起始值、范围的结束值以及每次迭代的步长值。

    2. 一个初始化(init)块,其中包含在 parfor 开始执行之前要执行一次的指令。

    3. 一个循环体,包含一组基本块,这些基本块对应于循环体,并计算迭代空间中的一个点。

    4. 用于迭代空间每个维度的索引变量。

    对于 pranges,循环嵌套是一个单一入口,其中起始、停止和步长字段来自指定的 prangeprange 的 parfor 的初始块为空,循环体是循环中的块集合减去循环头。

    启用并行化后,数组推导式(列表推导式)也将被翻译为 prange,以便并行运行。可以通过设置 parallel={'comprehension': False} 来禁用此行为。

    同样地,通过设置 parallel={'prange': False} 可以禁用从 prangeparfor 的转换,在这种情况下,prange 的处理方式与 range 相同。

  5. Numpy 到 parfor

    在这个子过程中,Numpy 函数如 oneszerosdot、大多数随机数生成函数、数组表达式(来自章节 阶段 5a:重写类型化 IR)以及 Numpy 归约操作被转换为 parfors。通常,这种转换会创建一个循环嵌套列表,其长度等于 IR 中赋值指令左侧的维数。左侧数组的维数和大小从上述子过程 3 中生成的数组分析信息中获取。生成一个创建结果 Numpy 数组的指令,并存储在新 parfor 的初始块中。为循环体创建一个基本块,并生成一个指令并添加到该块的末尾,以将计算结果存储到数组在迭代空间中的当前点。存储到数组中的结果取决于正在转换的操作。例如,对于 ones,存储的值是一个常数 1。对于生成随机数组的调用,值来自对相同随机数函数的调用,但删除了大小参数,因此返回一个标量。对于数组表达式操作符,数组表达式树被转换为 Numba IR,并且该表达式树根部的值被用于写入输出数组。可以通过设置 parallel={'numpy': False} 来禁用从 Numpy 函数和数组表达式操作符到 parfor 的转换。

    对于归约操作,循环嵌套列表同样使用数组分析信息为正在归约的数组创建。在初始化块中,初始值被赋给归约变量。循环体由一个单独的块组成,在该块中,迭代空间中的下一个值被获取,归约操作应用于该值和当前归约值,结果存储回归约值中。可以通过设置 parallel={'reduction': False} 来禁用将归约函数转换为 parfor

    NUMBA_DEBUG_ARRAY_OPT_STATS 环境变量设置为 1 将显示有关 parfor 转换的一些统计信息。

  6. 将项目设置为并行

    使用切片或布尔数组选择设置数组元素的范围也可以并行运行。如 A[P] = B[Q] (或更简单的 A[P] = c ,其中 c 是标量)的语句,如果满足以下条件之一,则会被翻译为 parfor

    1. PQ 是包含标量和切片的切片或多维选择器,并且 A[P]B[Q] 被数组分析认为是大小等价的。仅支持2值切片/范围,带有步长的3值切片将不会被翻译为 parfor

    2. PQ 是相同的布尔数组。

    可以通过设置 parallel={'setitem': False} 来禁用此翻译。

  7. 简化

    执行复制传播和死代码消除过程。

  8. 融合

    此子过程首先处理每个基本块,并在块内对指令进行重新排序,目的是将 parfors 推到块的下方,并将非 parfors 提升到块的开始部分。实际上,这种方法在 IR 中很好地将 parfors 彼此相邻放置,从而使得更多的 parfors 能够被融合。在 parfor 融合过程中,每个基本块会被反复扫描,直到无法再进行融合为止。在此扫描过程中,每一组相邻的指令都会被考虑。如果满足以下条件,相邻的指令将被融合在一起:

    1. 它们都是 parfors

    2. parfors 的循环嵌套大小相同,并且每个循环嵌套维度的数组等价类相同,并且

    3. 第一个 parfor 不会创建第二个 parfor 使用的归约变量。

    通过将第二个 parfor 的初始块添加到第一个 parfor 的初始块中,合并两个 parfor 的循环体,并用第一个 parfor 的循环索引变量替换第二个 parfor 体中的循环索引变量实例,将两个 parfor 融合在一起。可以通过设置 parallel={'fusion': False} 来禁用融合。

    NUMBA_DEBUG_ARRAY_OPT_STATS 环境变量设置为 1 将显示有关 parfor 融合的一些统计信息。

  9. 推送调用对象并计算parfor参数

    在第 阶段 6a:生成 nopython LLVM IR 节描述的降低阶段,每个 parfor 成为一个单独的函数,以 guvectorize (@guvectorize 装饰器) 风格并行执行。由于 parfor 可能使用函数中先前定义的变量,当这些 parfor 成为单独的函数时,这些变量必须作为参数传递给 parfor 函数。在这个子过程中,对每个 parfor 主体进行了一次使用-定义扫描,并使用活跃性信息来确定哪些变量被使用但未被 parfor 定义。该变量列表存储在此处的 parfor 中,以供降低阶段使用。函数变量在这个过程中是一个特殊情况,因为函数变量不能传递给以 nopython 模式编译的函数。相反,对于函数变量,此子过程将赋值指令推送到 parfor 主体中,以便这些变量不需要作为参数传递。

    要查看上述子阶段之间的中间IR以及其他调试信息,请将 NUMBA_DEBUG_ARRAY_OPT 环境变量设置为1。对于第 阶段 5a:重写类型化 IR 节中的示例,在此阶段会生成以下带有parfor的IR:

    ______________________________________________________________________
    label 0:
        a0 = arg(0, name=a0)                     ['a0']
        a0_sh_attr0.0 = getattr(attr=shape, value=a0) ['a0', 'a0_sh_attr0.0']
        $consta00.1 = const(int, 0)              ['$consta00.1']
        a0size0.2 = static_getitem(value=a0_sh_attr0.0, index_var=$consta00.1, index=0) ['$consta00.1', 'a0_sh_attr0.0', 'a0size0.2']
        a1 = arg(1, name=a1)                     ['a1']
        a1_sh_attr0.3 = getattr(attr=shape, value=a1) ['a1', 'a1_sh_attr0.3']
        $consta10.4 = const(int, 0)              ['$consta10.4']
        a1size0.5 = static_getitem(value=a1_sh_attr0.3, index_var=$consta10.4, index=0) ['$consta10.4', 'a1_sh_attr0.3', 'a1size0.5']
        a2 = arg(2, name=a2)                     ['a2']
        a2_sh_attr0.6 = getattr(attr=shape, value=a2) ['a2', 'a2_sh_attr0.6']
        $consta20.7 = const(int, 0)              ['$consta20.7']
        a2size0.8 = static_getitem(value=a2_sh_attr0.6, index_var=$consta20.7, index=0) ['$consta20.7', 'a2_sh_attr0.6', 'a2size0.8']
    ---begin parfor 0---
    index_var =  parfor_index.9
    LoopNest(index_variable=parfor_index.9, range=0,a0size0.2,1 correlation=5)
    init block:
        $np_g_var.10 = global(np: <module 'numpy' from '/usr/local/lib/python3.5/dist-packages/numpy/__init__.py'>) ['$np_g_var.10']
        $empty_attr_attr.11 = getattr(attr=empty, value=$np_g_var.10) ['$empty_attr_attr.11', '$np_g_var.10']
        $np_typ_var.12 = getattr(attr=float64, value=$np_g_var.10) ['$np_g_var.10', '$np_typ_var.12']
        $0.5 = call $empty_attr_attr.11(a0size0.2, $np_typ_var.12, kws=(), func=$empty_attr_attr.11, vararg=None, args=[Var(a0size0.2, test2.py (7)), Var($np_typ_var.12, test2.py (7))]) ['$0.5', '$empty_attr_attr.11', '$np_typ_var.12', 'a0size0.2']
    label 1:
        $arg_out_var.15 = getitem(value=a0, index=parfor_index.9) ['$arg_out_var.15', 'a0', 'parfor_index.9']
        $arg_out_var.16 = getitem(value=a1, index=parfor_index.9) ['$arg_out_var.16', 'a1', 'parfor_index.9']
        $arg_out_var.14 = $arg_out_var.15 * $arg_out_var.16 ['$arg_out_var.14', '$arg_out_var.15', '$arg_out_var.16']
        $arg_out_var.17 = getitem(value=a2, index=parfor_index.9) ['$arg_out_var.17', 'a2', 'parfor_index.9']
        $expr_out_var.13 = $arg_out_var.14 + $arg_out_var.17 ['$arg_out_var.14', '$arg_out_var.17', '$expr_out_var.13']
        $0.5[parfor_index.9] = $expr_out_var.13  ['$0.5', '$expr_out_var.13', 'parfor_index.9']
    ----end parfor 0----
        $0.6 = cast(value=$0.5)                  ['$0.5', '$0.6']
        return $0.6                              ['$0.6']
    ______________________________________________________________________
    

阶段 6a:生成 nopython LLVM IR

如果在类型推断过程中成功地为每个中间变量找到了Numba类型,那么Numba可以(潜在地)生成专门的本地代码。这个过程被称为 降低 。Numba IR树通过使用 llvmlite 中的辅助类被翻译成LLVM IR。机器生成的LLVM IR可能看起来过于冗长,但LLVM工具链能够相当容易地将其优化为紧凑、高效的代码。

基本的降低算法是通用的,但特定Numba IR节点如何转换为LLVM指令的具体细节由编译时选择的target context处理。默认的target context是“cpu”上下文,定义在 numba.targets.cpu 中。

可以通过将 NUMBA_DUMP_LLVM 环境变量设置为 1 来显示 LLVM IR。对于“cpu”上下文,我们的 add() 示例将如下所示:

define i32 @"__main__.add$1.int64.int64"(i64* %"retptr",
                                         {i8*, i32}** %"excinfo",
                                         i8* %"env",
                                         i64 %"arg.a", i64 %"arg.b")
{
   entry:
     %"a" = alloca i64
     %"b" = alloca i64
     %"$0.3" = alloca i64
     %"$0.4" = alloca i64
     br label %"B0"
   B0:
     store i64 %"arg.a", i64* %"a"
     store i64 %"arg.b", i64* %"b"
     %".8" = load i64* %"a"
     %".9" = load i64* %"b"
     %".10" = add i64 %".8", %".9"
     store i64 %".10", i64* %"$0.3"
     %".12" = load i64* %"$0.3"
     store i64 %".12", i64* %"$0.4"
     %".14" = load i64* %"$0.4"
     store i64 %".14", i64* %"retptr"
     ret i32 0
}

通过将 NUMBA_DUMP_OPTIMIZED 设置为 1,可以输出优化后的 LLVM IR。优化器显著缩短了上述生成的代码:

define i32 @"__main__.add$1.int64.int64"(i64* nocapture %retptr,
                                         { i8*, i32 }** nocapture readnone %excinfo,
                                         i8* nocapture readnone %env,
                                         i64 %arg.a, i64 %arg.b)
{
   entry:
     %.10 = add i64 %arg.b, %arg.a
     store i64 %.10, i64* %retptr, align 8
     ret i32 0
}

如果在 并行加速器 期间创建,parfor 操作将以以下方式降低。首先,parfor 的初始化块中的指令使用正常的降低代码降低到现有函数中。其次,parfor 的循环体被转换为单独的 GUFunc。第三,为当前函数生成调用并行 GUFunc 的代码。

要从 parfor 主体创建一个 GUFunc,GUFunc 的签名是通过采用阶段 阶段 5b:执行自动并行化 步骤 9 中识别的 parfor 参数,并向其添加一个特殊的 schedule 参数来创建的,GUFunc 将在此参数上并行化。schedule 参数实际上是一个静态调度,将 parfor 迭代空间的部分映射到 Numba 线程,因此 schedule 数组的长度与配置的 Numba 线程数相同。为了使这个过程更容易,并且对 Numba IR 的变化依赖性较小,此阶段创建了一个包含 GUFunc 参数和迭代代码的 Python 函数文本,该代码接受当前的 schedule 条目并循环遍历迭代空间中指定的部分。在该循环的主体中,插入了一个特殊的标记,以便后续容易定位。然后,处理迭代空间的代码通过 eval 生成,并调用 Numba 编译器的 run_frontend 函数来生成 IR。扫描该 IR 以定位标记,并用 parfor 的循环主体替换该标记。然后,通过使用 Numba 编译器的 compile_ir 函数编译这个合并的 IR,完成创建并行 GUFunc 的过程。

要调用并行 GUFunc,必须创建静态调度。插入代码以调用名为 do_scheduling 的函数。此函数以 parfor 的每个维度的尺寸和配置的 Numba 线程数 N`(:envvar:`NUMBA_NUM_THREADS)作为参数调用。do_scheduling 函数将迭代空间划分为 N 个大致相等大小的区域(1D 为线性,2D 为矩形,3D 及以上为超矩形),并将生成的调度传递给并行 GUFunc。分配给完整迭代空间某个维度的线程数大致与该维度大小与迭代空间所有维度大小总和的比率成正比。

并行归约不是 GUFuncs 原生提供的功能,但 parfor 降低策略允许我们以一种可以并行执行归约的方式使用 GUFuncs。为了实现这一点,对于 parfor 计算的每个归约变量,并行 GUFunc 和调用它的代码被修改,将标量归约变量转换为长度等于 Numba 线程数的归约变量数组。此外,GUFunc 仍然包含一个标量版本的归约变量,该变量在每次迭代期间由 parfor 主体更新。在 GUFunc 结束时的一次,这个本地归约变量被复制到归约数组中。通过这种方式,防止了归约数组的错误共享。在并行 GUFunc 返回后,主函数中还插入了代码,该代码对这个较小的归约数组进行归约,并将这个最终的归约值存储到原始的标量归约变量中。

对应于 阶段 5b:执行自动并行化 部分的示例的 GUFunc 如下所示:

______________________________________________________________________
label 0:
    sched.29 = arg(0, name=sched)            ['sched.29']
    a0 = arg(1, name=a0)                     ['a0']
    a1 = arg(2, name=a1)                     ['a1']
    a2 = arg(3, name=a2)                     ['a2']
    _0_5 = arg(4, name=_0_5)                 ['_0_5']
    $3.1.24 = global(range: <class 'range'>) ['$3.1.24']
    $const3.3.21 = const(int, 0)             ['$const3.3.21']
    $3.4.23 = getitem(value=sched.29, index=$const3.3.21) ['$3.4.23', '$const3.3.21', 'sched.29']
    $const3.6.28 = const(int, 1)             ['$const3.6.28']
    $3.7.27 = getitem(value=sched.29, index=$const3.6.28) ['$3.7.27', '$const3.6.28', 'sched.29']
    $const3.8.32 = const(int, 1)             ['$const3.8.32']
    $3.9.31 = $3.7.27 + $const3.8.32         ['$3.7.27', '$3.9.31', '$const3.8.32']
    $3.10.36 = call $3.1.24($3.4.23, $3.9.31, kws=[], func=$3.1.24, vararg=None, args=[Var($3.4.23, <string> (2)), Var($3.9.31, <string> (2))]) ['$3.1.24', '$3.10.36', '$3.4.23', '$3.9.31']
    $3.11.30 = getiter(value=$3.10.36)       ['$3.10.36', '$3.11.30']
    jump 1                                   []
label 1:
    $28.2.35 = iternext(value=$3.11.30)      ['$28.2.35', '$3.11.30']
    $28.3.25 = pair_first(value=$28.2.35)    ['$28.2.35', '$28.3.25']
    $28.4.40 = pair_second(value=$28.2.35)   ['$28.2.35', '$28.4.40']
    branch $28.4.40, 2, 3                    ['$28.4.40']
label 2:
    $arg_out_var.15 = getitem(value=a0, index=$28.3.25) ['$28.3.25', '$arg_out_var.15', 'a0']
    $arg_out_var.16 = getitem(value=a1, index=$28.3.25) ['$28.3.25', '$arg_out_var.16', 'a1']
    $arg_out_var.14 = $arg_out_var.15 * $arg_out_var.16 ['$arg_out_var.14', '$arg_out_var.15', '$arg_out_var.16']
    $arg_out_var.17 = getitem(value=a2, index=$28.3.25) ['$28.3.25', '$arg_out_var.17', 'a2']
    $expr_out_var.13 = $arg_out_var.14 + $arg_out_var.17 ['$arg_out_var.14', '$arg_out_var.17', '$expr_out_var.13']
    _0_5[$28.3.25] = $expr_out_var.13        ['$28.3.25', '$expr_out_var.13', '_0_5']
    jump 1                                   []
label 3:
    $const44.1.33 = const(NoneType, None)    ['$const44.1.33']
    $44.2.39 = cast(value=$const44.1.33)     ['$44.2.39', '$const44.1.33']
    return $44.2.39                          ['$44.2.39']
______________________________________________________________________

阶段 6b:生成对象模式 LLVM IR

如果类型推断无法为函数内的所有值找到 Numba 类型,该函数将以对象模式编译。生成的 LLVM 代码将显著变长,因为编译后的代码需要调用 Python C API 来执行基本上的所有操作。我们示例 add() 函数的优化 LLVM 代码如下:

@PyExc_SystemError = external global i8
@".const.Numba_internal_error:_object_mode_function_called_without_an_environment" = internal constant [73 x i8] c"Numba internal error: object mode function called without an environment\00"
@".const.name_'a'_is_not_defined" = internal constant [24 x i8] c"name 'a' is not defined\00"
@PyExc_NameError = external global i8
@".const.name_'b'_is_not_defined" = internal constant [24 x i8] c"name 'b' is not defined\00"

define i32 @"__main__.add$1.pyobject.pyobject"(i8** nocapture %retptr, { i8*, i32 }** nocapture readnone %excinfo, i8* readnone %env, i8* %arg.a, i8* %arg.b) {
entry:
  %.6 = icmp eq i8* %env, null
  br i1 %.6, label %entry.if, label %entry.endif, !prof !0

entry.if:                                         ; preds = %entry
  tail call void @PyErr_SetString(i8* @PyExc_SystemError, i8* getelementptr inbounds ([73 x i8]* @".const.Numba_internal_error:_object_mode_function_called_without_an_environment", i64 0, i64 0))
  ret i32 -1

entry.endif:                                      ; preds = %entry
  tail call void @Py_IncRef(i8* %arg.a)
  tail call void @Py_IncRef(i8* %arg.b)
  %.21 = icmp eq i8* %arg.a, null
  br i1 %.21, label %B0.if, label %B0.endif, !prof !0

B0.if:                                            ; preds = %entry.endif
  tail call void @PyErr_SetString(i8* @PyExc_NameError, i8* getelementptr inbounds ([24 x i8]* @".const.name_'a'_is_not_defined", i64 0, i64 0))
  tail call void @Py_DecRef(i8* null)
  tail call void @Py_DecRef(i8* %arg.b)
  ret i32 -1

B0.endif:                                         ; preds = %entry.endif
  %.30 = icmp eq i8* %arg.b, null
  br i1 %.30, label %B0.endif1, label %B0.endif1.1, !prof !0

B0.endif1:                                        ; preds = %B0.endif
  tail call void @PyErr_SetString(i8* @PyExc_NameError, i8* getelementptr inbounds ([24 x i8]* @".const.name_'b'_is_not_defined", i64 0, i64 0))
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_DecRef(i8* null)
  ret i32 -1

B0.endif1.1:                                      ; preds = %B0.endif
  %.38 = tail call i8* @PyNumber_Add(i8* %arg.a, i8* %arg.b)
  %.39 = icmp eq i8* %.38, null
  br i1 %.39, label %B0.endif1.1.if, label %B0.endif1.1.endif, !prof !0

B0.endif1.1.if:                                   ; preds = %B0.endif1.1
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_DecRef(i8* %arg.b)
  ret i32 -1

B0.endif1.1.endif:                                ; preds = %B0.endif1.1
  tail call void @Py_DecRef(i8* %arg.b)
  tail call void @Py_DecRef(i8* %arg.a)
  tail call void @Py_IncRef(i8* %.38)
  tail call void @Py_DecRef(i8* %.38)
  store i8* %.38, i8** %retptr, align 8
  ret i32 0
}

declare void @PyErr_SetString(i8*, i8*)

declare void @Py_IncRef(i8*)

declare void @Py_DecRef(i8*)

declare i8* @PyNumber_Add(i8*, i8*)

细心的读者可能会注意到在生成的代码中有几个不必要的 Py_IncRefPy_DecRef 调用。目前 Numba 还无法优化这些调用。

对象模式编译还会尝试识别可以提取并静态类型化为“nopython”编译的循环。这个过程称为*循环提升*,结果是创建一个仅包含循环的隐藏 nopython 模式函数,然后从原始函数中调用该函数。循环提升有助于提高需要访问不可编译代码(如 I/O 或绘图代码)但仍包含可编译代码的时间密集型部分的功能的性能。

第7阶段:将LLVM IR编译为机器代码

对象模式nopython 模式 中,生成的 LLVM IR 由 LLVM JIT 编译器编译,机器代码被加载到内存中。同时也会创建一个 Python 包装器(定义在 numba.dispatcher.Dispatcher 中),如果生成了多个类型的特化版本(例如,对于同一个函数的 float32float64 版本),它可以进行动态分派到正确的编译函数版本。

通过将 NUMBA_DUMP_ASSEMBLY 环境变量设置为 1,可以将 LLVM 生成的机器汇编代码转储到屏幕上:

        .globl  __main__.add$1.int64.int64
        .align  16, 0x90
        .type   __main__.add$1.int64.int64,@function
__main__.add$1.int64.int64:
        addq    %r8, %rcx
        movq    %rcx, (%rdi)
        xorl    %eax, %eax
        retq

汇编输出还将包括生成的包装函数,该函数将Python参数转换为本机数据类型。