使用内联其他函数编译带有函数的 Numba 模块时出错

Question

Numba documentation specifies that other compiled functions can be inlined and called from other compiled functions. This does not seem to be true when compiling ahead of time.

例如：这里有两个函数计算 2 个向量数组之间的内点积，其中一个计算实际乘积，另一个在循环内进行内联调用：

# Module test.py
import numpy as np
from numba import njit, float64

@njit(float64(float64[:], float64[:]))
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

@njit(float64[:](float64[:,:], float64[:,:]))
def n_inner1d(a, b):
    prod = np.empty(a.shape[0])    
    for i in range(a.shape[0]):
        prod[i] = product(a[i], b[i])

    return prod

照原样，我可以 import test 并使用 test.n_inner1d 非常好。现在让我们做一些修改，以便可以将其编译为 .pyd

# Module test.py
import numpy as np
from numba import float64
from numba.pycc import CC

cc = CC('test')
cc.verbose = True

@cc.export('product','float64(float64[:], float64[:])')
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

@cc.export('n_inner1d','float64[:](float64[:,:], float64[:,:])')
def n_inner1d(a, b):
    prod = np.empty(a.shape[0])    
    for i in range(a.shape[0]):
        prod[i] = product(a[i], b[i])

    return prod

if __name__ == "__main__":
    cc.compile()

尝试编译时，出现以下错误：

# python test.py
Failed at nopython (nopython frontend)
Untyped global name 'product': cannot determine Numba type of <type 'function'>
File "test.py", line 20

问题

对于已编译的模块 ahead of time，其中定义的函数是否可以相互调用并内联使用？

Answer 1

我联系了 numba 开发人员，他们友好地回答说在 @cc.export 之后添加 @njit 装饰器将使函数调用类型解析工作并解析。

例如：

@cc.export('product','float64(float64[:], float64[:])')
@njit
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

将使 product 功能对其他人可用。需要注意的是，在某些情况下，内联函数完全有可能以与声明的 AOT 不同的类型签名结束。

使用内联其他函数编译带有函数的 Numba 模块时出错

Error when compiling a Numba module with function using other functions inline

python

numpy

numba

问题