Numba AOT编译含函数参数的函数时如何正确指定参数签名

问题原因 #

@cc.export如果传入字符串形式的签名，字符串内部的nb.typeof这类Python表达式不会被执行解析，仅会被当作普通的类型名字符串处理，自然无法识别nb、fcn_sum这类变量，就会触发NameError。

正确实现方案 #

不要使用字符串形式定义带函数参数的签名，直接通过Numba类型对象构建签名传入即可，完整示例代码如下：

import numba as nb
from numba.pycc import CC

# 初始化AOT编译模块
cc = CC('test_aot_module')

# 先编译作为参数传递的目标函数
@nb.njit(nb.f8(nb.f8, nb.f8))
def fcn_sum(a, b): 
    return a + b

# 获取对应函数的类型对象
func_param_type = nb.typeof(fcn_sum)
# 直接传入构造好的签名对象，不要用字符串形式
@cc.export('test', nb.f8(func_param_type, nb.f8, nb.f8))
def test(fun, a, b):
    return fun(a, b)

# 执行编译生成pyd/so扩展文件
if __name__ == "__main__":
    cc.compile()

注意事项 #

• AOT属于静态提前编译，编译时就固定了参数类型，后续调用test时传入的函数参数，必须是签名为f8(f8, f8)的Numba jitted函数，不支持动态匹配其他签名的函数。
• 只要其他jitted函数的输入输出签名和fcn_sum完全一致，也可以作为参数传入该AOT编译后的test函数使用。

内容的提问来源于stack exchange，提问作者7vn_