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 使用基于栈的解释器(很像惠普计算器),因此代码首先将两个局部变量压入栈中。BINARY_ADD
操作码将栈顶的两个参数弹出,并进行一个 Python C API 函数调用,该调用等同于调用 a.__add__(b)
。结果随后被压入解释器栈顶。最后,RETURN_VALUE
操作码返回栈顶的值作为函数调用的结果。
Numba 可以获取此字节码并将其编译为机器码,执行与 CPython 解释器相同的操作,将 a
和 b
视为通用 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”和“cuda”上下文,以及一个生成多线程 CPU 代码的“parallel”上下文。
编译器阶段
Numba 中的 jit()
装饰器最终调用 numba.compiler.compile_extra()
,该函数通过一个多阶段过程编译 Python 函数,如下所述。
阶段 1:分析字节码
编译开始时,函数字节码被传递给 Numba 解释器 (numba.interpreter
) 的一个实例。解释器对象分析字节码以找到控制流图 (numba.controlflow
)。控制流图 (CFG) 描述了由于循环和分支,执行如何在函数内部从一个块移动到下一个块。
数据流分析 (numba.dataflow
) 接收控制流图并追踪值如何在 Python 解释器栈上被压入和弹出,以用于不同的代码路径。这对于理解栈上变量的生命周期很重要,而这些在阶段 2 中是必需的。
如果您将环境变量 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,对象模式回退只会在启用循环提升时发生。
阶段 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 只对数据进行一次遍历,从而通常具有更好的缓存性能。最后,在降级过程中,这些 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。子阶段在以下段落中进行了更详细的描述。
- CFG 简化
有时 Numba IR 会包含不含循环的块链,此子阶段会将这些块合并为单个块。此子阶段简化了后续的 IR 分析。
- Numpy 规范化
某些 Numpy 操作可以写为对 Numpy 对象的操作(例如
arr.sum()
),或写为接受这些对象的 Numpy 调用(例如numpy.sum(arr)
)。此子阶段将所有此类操作转换为后一种形式,以便进行更清晰的后续分析。
- 数组分析
后续 parfor 融合的一个关键要求是 parfor 具有相同的迭代空间,并且这些迭代空间通常对应于 Numpy 数组维度的大小。在此子阶段中,IR 被分析以确定 Numpy 数组维度的等价类。考虑示例
a = b + 1
,其中a
和b
都是 Numpy 数组。这里,我们知道a
的每个维度必须与b
的相应维度具有相同的等价类。通常,富含 Numpy 操作的例程将使得函数中创建的所有数组的等价类完全可知。数组分析还将推断切片选择和布尔数组掩码(仅一维)的大小等效性。例如,它能够推断
a[1 : n-1]
的大小与b[0 : n-2]
相同。数组分析还可能插入安全假设,以确保在操作并行化之前满足与数组大小相关的预条件。例如,2D 矩阵
X
和 1D 向量w
之间的np.dot(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
明确的断言有助于消除函数其余部分中的所有边界检查。
prange()
到 parfor在 for 循环中使用 prange (显式并行循环) 是程序员的明确指示,表明 for 循环的所有迭代都可以并行执行。在此子阶段中,我们分析 CFG 以定位循环,并将那些由 prange 对象控制的循环转换为显式 parfor 操作符。每个显式 parfor 操作符由以下部分组成:
描述 parfor 迭代空间的循环嵌套信息列表。循环嵌套列表中的每个条目包含一个索引变量、范围的起始、范围的结束以及每次迭代的步长值。
一个初始化(init)块,其中包含在 parfor 开始执行之前执行一次的指令。
一个循环体,由一组基本块组成,这些基本块对应于循环的主体,并计算迭代空间中的一个点。
用于迭代空间每个维度的索引变量。
对于 parfor pranges,循环嵌套是单个条目,其中起始、停止和步长字段来自指定的 prange。对于 prange parfor,init 块是空的,循环体是循环中除了循环头之外的块集合。
启用并行化后,数组推导式(列表推导式)也将转换为 prange 以便并行运行。可以通过设置
parallel={'comprehension': False}
来禁用此行为。同样,可以通过设置
parallel={'prange': False}
来禁用整体的 prange 到 parfor 转换,在这种情况下,prange 将被视为与 range 相同。
- Numpy 到 parfor
在此子阶段中,Numpy 函数(如
ones
、zeros
、dot
)、大多数随机数生成函数、数组表达式(来自阶段 5a:重写类型化 IR)以及 Numpy 归约都被转换为 parfors。通常,此转换会创建循环嵌套列表,其长度等于 IR 中赋值指令左侧的维数。左侧数组的维数和大小取自上面子阶段 3 中生成的数组分析信息。生成创建结果 Numpy 数组的指令,并存储在新 parfor 的 init 块中。为循环体创建一个基本块,并生成一条指令并添加到该块的末尾,以将计算结果存储到迭代空间当前点的数组中。存储到数组中的结果取决于正在转换的操作。例如,对于ones
,存储的值是常量 1。对于生成随机数组的调用,该值来自对相同随机数函数的调用,但删除了 size 参数,因此返回一个标量。对于 arrayexpr 操作符,arrayexpr 树被转换为 Numba IR,并且该表达式树根部的值用于写入输出数组。通过设置parallel={'numpy': False}
可以禁用 Numpy 函数和 arrayexpr 操作符到 parfor 的转换。对于归约操作,循环嵌套列表也类似地使用被归约数组的数组分析信息创建。在 init 块中,初始值被赋值给归约变量。循环体由一个单块组成,在该块中,迭代空间中的下一个值被取出,然后将归约操作应用于该值和当前归约值,并将结果存回归约变量。通过设置
parallel={'reduction': False}
可以禁用归约函数到 parfor 的转换。将
NUMBA_DEBUG_ARRAY_OPT_STATS
环境变量设置为 1 将显示关于 parfor 转换的一些统计信息。
- Setitem 到 parfor
使用切片或布尔数组选择来设置数组元素范围也可以并行运行。如果满足以下条件之一,语句如
A[P] = B[Q]
(或更简单的情况A[P] = c
,其中c
是一个标量)将被转换为 parfor:P
和Q
是切片或涉及标量和切片的多维选择器,并且A[P]
和B[Q]
被数组分析认为是大小等价的。仅支持 2 值切片/范围,带步长的 3 值切片将不会被转换为 parfor。P
和Q
是相同的布尔数组。
通过设置
parallel={'setitem': False}
可以禁用此转换。
- 简化
执行一次副本传播和死代码消除阶段。
- 融合
此子阶段首先处理每个基本块,并对块内的指令进行重新排序,目标是将 parfors 推到块的下部,并将非 parfors 提升到块的开头。在实践中,这种方法很好地使 parfors 在 IR 中彼此相邻,从而能够融合更多的 parfors。在 parfor 融合期间,每个基本块都被反复扫描,直到无法进一步融合。在此扫描期间,会考虑每组相邻的指令。相邻指令在以下情况下融合:
它们都是 parfors
parfors 的循环嵌套大小相同,并且循环嵌套的每个维度的数组等价类也相同,以及
第一个 parfor 不创建被第二个 parfor 使用的归约变量。
两个 parfor 通过将第二个 parfor 的 init 块添加到第一个 parfor 的 init 块中,合并两个 parfor 的循环体,并将第二个 parfor 循环体中第二个 parfor 的循环索引变量实例替换为第一个 parfor 的循环索引变量来融合。可以通过设置
parallel={'fusion': False}
来禁用融合。将
NUMBA_DEBUG_ARRAY_OPT_STATS
环境变量设置为 1 将显示关于 parfor 融合的一些统计信息。
- 推送调用对象并计算 parfor 参数
在阶段 6a:生成 nopython LLVM IR 中描述的降级阶段,每个 parfor 都会成为一个单独的函数,以
guvectorize
(@guvectorize 装饰器)风格并行执行。由于 parfors 可能会使用函数中先前定义的变量,当这些 parfors 成为单独的函数时,这些变量必须作为参数传递给 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 指令的具体细节由为编译选择的目标上下文处理。默认的目标上下文是“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
}
优化后的 LLVM IR 可以通过将 NUMBA_DUMP_OPTIMIZED
设置为 1 来输出。优化器大大缩短了上面生成的代码
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
}
如果在阶段 5b:执行自动并行化期间创建,parfor 操作会以下列方式降级。首先,parfor 的 init 块中的指令使用正常的降级代码降级到现有函数中。其次,parfor 的循环体转换为单独的 GUFunc。第三,为当前函数发出代码以调用并行 GUFunc。
要从 parfor 主体创建 GUFunc,首先根据阶段 5b:执行自动并行化的步骤 9 中识别出的 parfor 参数来创建 GUFunc 的签名,并添加一个特殊的 schedule 参数,GUFunc 将通过该参数进行并行化。实际上,schedule 参数是一个静态调度,将 parfor 迭代空间的部分映射到 Numba 线程,因此 schedule 数组的长度与配置的 Numba 线程数相同。为了使这个过程更简单,并且减少对 Numba IR 更改的依赖,此阶段创建了一个包含 GUFunc 参数和迭代代码的 Python 函数文本,该迭代代码获取当前调度条目并遍历迭代空间的指定部分。在该循环体中,插入了一个特殊的哨兵,以便后续轻松定位。然后,处理迭代空间的这段代码被 eval
执行,并调用 Numba 编译器的 run_frontend 函数来生成 IR。扫描该 IR 以定位哨兵,并将哨兵替换为 parfor 的循环体。最后,通过使用 Numba 编译器的 compile_ir
函数编译这个合并的 IR,完成创建并行 GUFunc 的过程。
要调用并行 GUFunc,必须创建静态调度。插入代码以调用名为 do_scheduling.
的函数。此函数以每个 parfor 维度的大小以及配置的 Numba 线程数 N(NUMBA_NUM_THREADS
)作为参数被调用。do_scheduling
函数将迭代空间划分为 N 个近似相等大小的区域(1D 为线性,2D 为矩形,3D 以上为超矩形),并将生成的调度传递给并行 GUFunc。分配给完整迭代空间给定维度的线程数大致与给定维度的大小与迭代空间所有维度大小之和的比率成比例。
GUFunc 本身不原生提供并行归约功能,但 parfor 降级策略允许我们以并行方式执行归约操作。为此,对于 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_IncRef
和 Py_DecRef
的几次不必要的调用。目前 Numba 无法优化掉这些调用。
对象模式编译还会尝试识别可以提取并静态类型化以进行“nopython”编译的循环。这个过程被称为*循环提升*,它会创建一个只包含循环的隐藏 nopython 模式函数,然后从原始函数中调用该函数。循环提升有助于提高需要访问不可编译代码(例如 I/O 或绘图代码)但仍包含时间密集型可编译代码部分的函数的性能。
阶段 7:将 LLVM IR 编译为机器码
在对象模式和nopython 模式下,生成的 LLVM IR 都由 LLVM JIT 编译器编译,机器码加载到内存中。同时还会创建一个 Python 包装器(定义在 numba.dispatcher.Dispatcher
中),如果生成了多个类型特化版本(例如,同一函数的 float32
和 float64
版本),该包装器可以动态分派到正确编译的函数版本。
通过将 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 参数转换为本地数据类型。