使用 numba 运行嵌套函数

Question

我最近尝试使用 numba 来加速我在 python 中的部分代码。我试图从函数 2 内部运行函数 1，而它们都是用 numba 编译的，但它不起作用。这是我的代码：

import numba as nb
from math import acos
from time import time

@nb.jit("void()")
def myfunc():
    s = 0
    for i in range(10000000):
        s += acos(0.5)
    print('The sum is: ', s)


@nb.jit("void()")
def myfunc2():
    myfunc()


tic = time()
myfunc2()
toc = time()
print(toc-tic)

当我调用 myfunc() 时，代码有效，并且我得到的结果比我不使用 numba 时快得多。但是，当我调用myfunc2 时，我看到了这个错误：

 File "~/.spyder-py3/temp.py", line 22, in <module>
    myfunc2()

RuntimeError: missing Environment

任何人都知道为什么在这种情况下从另一个 insdie 调用一个函数不起作用？

Answer 1

Numba v0.39+

在 v0.39 中引入了一个修复程序。根据Release Notes：

PR #2986：修复环境传播

更多详情请见github pull #2986。

Numba pre-v0.39

这是一个已知问题。如github issue #2411中所述：

似乎环境指针没有正确传递 nopython 函数。

如下修改以从 numba 函数中删除 print() 应该可以解决此问题：

import numba as nb
from math import acos
from time import time

@nb.jit("void()")
def myfunc():
    s = 0
    for i in range(10000000):
        s += acos(0.5)
    return s

@nb.jit("void()")
def myfunc2():
    return myfunc()

tic = time()
x = myfunc2()  # 10471975.511390356
toc = time()
print(toc-tic)