运行使用 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

有人知道为什么在这种情况下从独立调用另一个函数不起作用吗？

Answer 1

Numba v0.39+

在 v0.39 中引入了修复。根据 Release Notes:

PR #2986: Fix environment propagation

有关详细信息，请参阅 github pull #2986。

Numba v0.39 之前的版本

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

Seems like the environment pointer is not passed properly across nopython functions.

如下修改以从 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)

运行 使用 numba 的嵌套函数

Running nested functions using numba

python

jit

numba

Numba v0.39+

Numba v0.39 之前的版本

运行使用 numba 的嵌套函数