原文: http://numba.pydata.org/numba-doc/latest/developer/architecture.html
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
操作码从堆栈中弹出前两个参数,并进行 Python C API 函数调用,相当于调用a.__add__(b)
。然后将结果推送到解释器堆栈的顶部。最后,RETURN_VALUE
操作码作为函数调用的结果返回堆栈顶部的值。
Numba 可以使用此字节码并将其编译为执行与 CPython 解释器相同操作的机器代码,将a
和b
视为通用 Python 对象。保留了 Python 的完整语义,并且编译的函数可以与任何具有 add 运算符定义的对象一起使用。当以这种方式编译 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 代码的“并行”上下文。
Numba 中的 jit()
装饰器最终调用numba.compiler.compile_extra()
,它在多阶段过程中编译 Python 函数,如下所述。
在编译开始时,函数字节码被传递给 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 中的数字指的是上面操作码名称左侧显示的字节码偏移量。
一旦控制流和数据分析完成,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 的拥有引用。
现在该函数已被转换为 Numba IR,可以执行宏扩展。宏扩展将 Numba 已知的特定属性转换为表示函数调用的 IR 节点。这在numba.compiler.translate_stage
功能中启动,并在numba.macro
中实现。
宏扩展的属性示例包括网格,块和线程维度和索引的 CUDA 内在函数。例如,在以下函数中分配给tx
:
@cuda.jit(argtypes=[f4[:]])
def f(a):
tx = cuda.threadIdx.x
在翻译成 Numba IR 之后有以下代表:
$0.1 = global(cuda: <module 'numba.cuda' from '...'>) ['$0.1']
$0.2 = getattr(value=$0.1, attr=threadIdx) ['$0.1', '$0.2']
del $0.1 []
$0.3 = getattr(value=$0.2, attr=x) ['$0.2', '$0.3']
del $0.2 []
tx = $0.3 ['$0.3', 'tx']
宏扩展后,$0.3 = getattr(value=$0.2, attr=x)
IR 节点转换为:
$0.3 = call tid.x(, ) ['$0.3']
它表示用于调用tid.x
内部函数的Intrinsic
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') []
现在已经生成了 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 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 的函数,该函数内联到一个仅分配单个结果数组的循环中。
仅当 jit()
装饰器中的parallel
选项设置为True
时,才会执行此过程。此过程在 Numba IR 中的操作语义中隐含发现并行性,并使用特殊的 <cite>parfor</cite> 运算符替换那些操作的显式并行表示。然后,执行优化以最大化彼此相邻的 parfors 的数量,使得它们可以融合在一起,只需要一次通过数据,因此通常具有更好的高速缓存性能。最后,在降低期间,这些 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
}
对于所有这些,默认设置为 <cite>True</cite> 。在以下段落中更详细地描述了子通道。
-
CFG Simplification
有时 Numba IR 将包含不包含循环的块链,这些循环在此子过程中合并为单个块。该子通过简化了 IR 的后续分析。
-
Numpy canonicalization
一些 Numpy 操作可以写为 Numpy 对象上的操作(例如
arr.sum()
),或者作为 Numpy 对这些对象的调用(例如numpy.sum(arr)
)。该子过程将所有这些操作转换为后一种形式,以便进行更清晰的后续分析。 -
Array analysis
后期 parfor 融合的一个关键要求是 parfors 具有相同的迭代空间,这些迭代空间通常对应于 Numpy 数组的维度大小。在该子过程中,分析 IR 以确定 Numpy 阵列的维度的等价类。考虑示例
a = b + 1
,其中a
和b
都是 Numpy 数组。在这里,我们知道a
的每个维度必须与b
的相应维度具有相同的等价类。通常,富含 Numpy 操作的例程将使函数中创建的所有数组都能完全知道等价类。数组分析还将推断切片选择的大小等效性和布尔数组掩蔽(仅一维)。例如,它能够推断
a[1 : n-1]
与b[0 : n-2]
的大小相同。阵列分析还可以插入安全假设,以确保在并行操作之前满足与阵列大小相关的前提条件。例如,2-D 矩阵
X
和 1-D 矢量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() to parfor
在 for 循环中使用 prange(显式并行循环)是程序员明确表示 for 循环的所有迭代都可以并行执行。在这个子过程中,我们分析 CFG 以定位循环并将由 prange 对象控制的循环转换为显式 <cite>parfor</cite> 运算符。每个显式 parfor 运算符包括:
- 循环嵌套信息列表,用于描述 parfor 的迭代空间。循环嵌套列表中的每个条目都包含索引变量,范围的起点,范围的结束以及每次迭代的步长值。
- 初始化(init)块,包含在 parfor 开始执行之前执行一次的指令。
- 循环体包括一组基本块,其对应于循环体并计算迭代空间中的一个点。
- 用于迭代空间的每个维度的索引变量。
对于 parfor <cite>pranges</cite> ,循环嵌套是一个单独的条目,其中 start,stop 和 step 字段来自指定的 <cite>prange</cite> 。对于 <cite>prange</cite> parfors,init 块为空,循环体是循环中的块集减去循环头。
通过并行化,数组理解( List comprehension )也将被转换为 prange 以便并行运行。通过设置
parallel={'comprehension': False}
禁用此行为。同样,通过设置
parallel={'prange': False}
可以禁用整体 <cite>prange</cite> 到 <cite>parfor</cite> 转换,在这种情况下, <cite>prange</cite> 的处理方式与<cite>范围相同</cite> ]。 -
Numpy to parfor
在这个子通道中,Numpy 函数如
ones
,zeros
,dot
,大多数随机数生成函数,arrayexprs(来自 Section Stage 6a:Rewrite typed IR )和 Numpy 减少量转换为 parfors。通常,此转换会创建循环嵌套列表,其长度等于 IR 中赋值指令左侧的维数。左侧阵列的尺寸的数量和大小取自上面的子通道 3 中生成的阵列分析信息。创建结果 Numpy 数组的指令生成并存储在新 parfor 的 init 块中。为循环体创建基本块,并生成指令并将其添加到该块的末尾,以将计算结果存储到迭代空间中当前点的数组中。存储到数组中的结果取决于要转换的操作。例如,对于ones
,存储的值是常量 1.对于生成随机数组的调用,该值来自对相同随机数函数的调用,但是大小参数被删除,因此返回标量。对于 arrayexpr 运算符,arrayexpr 树将转换为 Numba IR,并且该表达式树的根处的值用于写入输出数组。可以通过设置parallel={'numpy': False}
禁用从 Numpy 函数和 arrayexpr 运算符到 <cite>parfor</cite> 的转换。对于缩减,类似地使用要减少的阵列的阵列分析信息来创建循环嵌套列表。在 init 块中,初始值被分配给 reduce 变量。循环体由单个块组成,其中取出迭代空间中的下一个值,并将缩减操作应用于该值,并将当前缩减值和结果存储回缩减值。通过设置
parallel={'reduction': False}
可以禁用将缩小功能转换为 <cite>parfor</cite> 。将
NUMBA_DEBUG_ARRAY_OPT_STATS
环境变量设置为 1 将显示有关 parfor 转换的一般统计信息。 -
Setitem to parfor
使用切片或布尔数组选择设置数组元素的范围也可以并行运行。如果满足下列条件之一,
A[P] = B[Q]
(或更简单的情况A[P] = c
,其中c
是标量)等语句将转换为 <cite>parfor</cite> :> 1.
P
和Q
是涉及标量和切片的切片或多维选择器,并且通过阵列分析将A[P]
和B[Q]
视为大小等效。仅支持 2 值切片/范围,带步长的 3 值不会转换为 <cite>parfor</cite> 。 > 2.P
和Q
是相同的布尔数组。可以通过设置
parallel={'setitem': False}
禁用此转换。 -
Simplification
执行复制传播和死代码消除传递。
-
Fusion
该子通道首先处理每个基本块并对块内的指令进行重新排序,目标是在块中将 parfors 推低,并将非 parfors 提升到块的开始。在实践中,这种方法很好地使得在 IR 中彼此相邻的 parfors,这使得更多的 parfors 可以融合。在 parfor 融合期间,重复扫描每个基本块直到不可能进一步融合。在此扫描期间,考虑每组相邻指令。如果符合以下条件,相邻指令将融合
- 他们都是 parfors
- parfors 的循环嵌套具有相同的大小,并且循环嵌套的每个维度的数组等价类是相同的,并且
- 第一个 parfor 不会创建第二个 parfor 使用的缩减变量。
将两个 parfors 融合在一起,将第二个 parfor 的 init 块添加到第一个,将两个 parfors 的循环体合并在一起,并用第一个 parfor 的循环索引变量替换第二个 parfor 的循环中第二个 parfor 的循环索引变量的实例。可以通过设置
parallel={'fusion': False}
来禁用 Fusion。将
NUMBA_DEBUG_ARRAY_OPT_STATS
环境变量设置为 1 将显示有关 parfor 融合的一些统计信息。 -
Push call objects and compute parfor parameters
在阶段 7a:生成 nopython LLVM IR 中描述的降低阶段中,每个 parfor 成为在
guvectorize
( @guvectorize 装饰器)样式中并行执行的单独函数。由于 parfors 可能使用先前在函数中定义的变量,当这些 parfors 成为单独的函数时,这些变量必须作为参数传递给 parfor 函数。在该子通道中,对每个 parfor 主体进行 use-def 扫描,并且活跃度信息用于确定使用但未由 parfor 定义的变量。该变量列表存储在 parfor 中,以便在降低期间使用。函数变量是此过程中的一个特例,因为函数变量无法传递给在 nopython 模式下编译的函数。相反,对于函数变量,此子过程将赋值指令推送到 parfor 主体中,以便不需要将它们作为参数传递。要查看上述子通道与其他调试信息之间的中间 IR,请将
NUMBA_DEBUG_ARRAY_OPT
环境变量设置为 1.对于第 6a 阶段:重写类型 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'] ______________________________________________________________________
如果类型推断成功为每个中间变量找到 Numba 类型,那么 Numba 可以(可能)生成专门的本机代码。该过程称为降低。通过使用来自 llvmlite 的辅助类,将 Numba IR 树转换为 LLVM IR。机器生成的 LLVM IR 看起来不必要地冗长,但 LLVM 工具链能够很容易地将其优化为紧凑,高效的代码。
基本降低算法是通用的,但特定 Numba IR 节点如何转换为 LLVM 指令的细节由选择用于编译的目标上下文处理。默认目标上下文是在numba.targets.cpu
中定义的“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
}
如果在阶段 6b:执行自动并行化期间创建,则以下列方式降低 parfor 操作。首先,使用正常的降低代码将 parfor 的 init 块中的指令降低到现有函数中。其次,parfor 的循环体转变为单独的 GUFunc。第三,为当前函数发出代码以调用并行 GUFunc。
要从 parfor 主体创建 GUFunc,GUFunc 的签名是通过获取 Stage Stage 6b 的步骤 9 中确定的 parfor 参数创建的:执行自动并行化并添加一个特殊的<cite>时间表</cite>参数,GUFunc 将在其中并行化。 schedule 参数实际上是将 parfor 迭代空间的部分映射到 Numba 线程的静态调度,因此调度数组的长度与配置的 Numba 线程的数量相同。为了使这个过程更容易并且更少依赖于对 Numba IR 的更改,此阶段创建一个 Python 函数作为文本,其中包含 GUFunc 的参数和迭代代码,该代码获取当前调度条目并循环遍历迭代空间的指定部分。在该循环的主体中,插入一个特殊的标记,以便随后轻松定位。处理迭代空间处理的代码然后被eval
编辑,并且调用 Numba 编译器的 run_frontend 函数来生成 IR。扫描 IR 以定位哨兵,并用穹顶的环体替换哨兵。然后,通过使用 Numba 编译器的compile_ir
函数编译此合并的 IR 来完成创建并行 GUFunc 的过程。
要调用并行 GUFunc,必须创建静态调度。插入代码以调用名为do_scheduling.
的函数。使用每个 parfor 维度的大小和配置的 Numba 线程的数量 <cite>N</cite> ( NUMBA_NUM_THREADS
)调用此函数。 do_scheduling
函数将迭代空间划分为 N 个近似相等大小的区域(1D 为线性,2D 为矩形,3 + D 为超矩形),结果时间表传递给并行 GUFunc。专用于完整迭代空间的给定维度的线程数大致与给定维度的大小与迭代空间的所有维度的大小之和的比率成比例。
GUFuncs 本身并不提供并行缩减,但降低 parfor 的策略允许我们以可以并行执行缩减的方式使用 GUFunc。为了实现这一点,对于由 parfor 计算的每个缩减变量,并行 GUFunc 和调用它的代码被修改以使标量缩减变量成为一个缩减变量数组,其长度等于 Numba 线程的数量。此外,GUFunc 仍然包含缩减变量的标量版本,该变量在每次迭代期间由 parfor 正文更新。一次在 GUFunc 结束时,这个局部缩减变量被复制到缩小数组中。以这种方式,防止了减少阵列的错误共享。在并行 GUFunc 返回后,代码也会插入到 main 函数中,该函数在这个较小的缩减数组中进行缩减,然后将此最终缩减值存储到原始标量缩减变量中。
对应于阶段 6b:执行自动并行化中的示例的 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']
______________________________________________________________________
如果类型推断无法为函数内的所有值找到 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 或绘图代码)的函数的性能,但仍包含可编译代码的时间密集部分。
在对象模式和 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 参数转换为本机数据类型。