numba 渴望编译？什么是模式？

Question

我在 numba 的网站上查看了 eager compilation，但不知道如何指定类型：

他们使用的例子是这样的：

from numba import jit, int32

@jit(int32(int32, int32))
def f(x, y):
    # A somewhat trivial example
    return x + y

# source: http://numba.pydata.org/numba-doc/latest/user/jit.html#eager-compilation

如您所见，它有 2 个变量作为输入，returns 有一个变量。它们都应该是 int32.

理解装饰器的一种方法是 @jit(int32(int32, int32)) 可以理解为：

@jit(type_of_returned_value(type_of_x, type_of_b))

如果那是对的（是吗？），那么如何为多个输入和输出指定它？

像这样：

@nb.jit
def filter3(a,b):
    return a > b 

@nb.jit
def func3(list_of_arrays_A, list_of_arrays_B, list_of_arrays_C, list_of_arrays_D, 2d_numpy_array_of_objects):

    for i in range(len(list_of_arrays_A)): 

        for j in range(list_of_arrays_A[i].size):
            
            if filter3(list_of_arrays_A[i][j],list_of_arrays_B[i][j]):
                2d_numpy_array_of_objects[i][j] = 1

            elif filter3(list_of_arrays_B[i][j],list_of_arrays_A[i][j]):
                2d_numpy_array_of_objects[i][j] = 0

            elif filter3(list_of_arrays_C[i][j],list_of_arrays_D[i][j]): 
                2d_numpy_array_of_objects[i][j] = 0
            else:                       
                2d_numpy_array_of_objects[i][j] = 1 
'''

My intention: Since i need to speed up a function which is only called once, (but takes forever if **not** done with numba), I need to speed up its numba-compilation

Answer 1

总是可以使用 numba.typeof 来干扰变量的类型。例如

import numpy as np
import numba as nb
N=10000
simple_list= [np.zeros(1) for x in range(N)]
nb.typeof(simple_list)
# reflected list(array(float64, 1d, C))

或：

from numba.typed import List
typed_list=List()
for _ in range(N):
    typed_list.append(np.zeros(1))
nb.typeof(typed_list)
# ListType[array(float64, 1d, C)]

所以你可以提供如下签名给ahead-of-time编译：

@nb.jit([nb.void(nb.typeof(typed_list)),
         nb.void(nb.typeof(simple_list))])
def fun(lst):
    pass

值得注意的细节：

我提前编译了两个不同版本的函数：一个用于 numba 的 TypedList (nb.void(nb.typeof(typed_list))，另一个用于 python 的列表 (nb.void(nb.typeof(simple_list)))。
我不使用签名字符串，而是使用签名本身（例如 here 中描述的），因为如果我理解正确的话 TypedList 或反射列表不存在签名字符串（更多信息如下所示）。
由于函数 fun 没有 return 任何东西，函数的 return 类型是 void，因此在签名中 nb.void(...)。

然而，有趣的是 simple_list-version 有多少开销：

%timeit fun(simple_list)  # 185 ms ± 4.23 ms 
%timeit fun(typed_list)   # 1.18 µs ± 69.3 ns

即约 1e5 的因数！原因也很清楚：为了检查传递的列表是否确实是 reflected list(array(float64, 1d, C)) 类型，numba 必须查看列表中的每个元素。另一方面，对于 TypedList 它要简单得多：列表中不能只有一种类型——不需要遍历整个列表！

因此，人们应该更愿意创建和使用 TypedList，而不仅仅是关闭 deprecation warning。

可能无法给出reflected list或TypedList的字符串，因为现在the following code用于解析签名：

def _parse_signature_string(signature_str):
    # Just eval signature_str using the types submodules as globals
    return eval(signature_str, {}, types.__dict__)

并且因为 nb.types.__dict__ 没有 TypedList 或 reflected list 我们不能通过字符串传递它们。

一旦函数被编译（ahead-of-time 或 just-in-time），就可以在相应的 Dispatcher-object 中看到签名，例如通过：

[x.signature for x in fun.overloads.values()]
# [(ListType[array(float64, 1d, C)],) -> none,
#  (reflected list(array(float64, 2d, C)),) -> none]

这个可以用来算出函数的右边return-type（这里的none就是void）。

numba 渴望编译？什么是模式？

numba eager compilation? Whats the pattern?

python

numpy

compilation

vectorization

numba