python - numba.jit (nopython) 可以处理数组索引和切片吗？

Question

我正在尝试使用 numba.jit 加速一个简单的 Python 循环。但似乎 jit 无法处理基本的数组索引和切片？我能做些什么来使这成为可能吗？如果 jit 不能处理基本的 numpy 数组，我不明白它怎么会有用。

我强制使用 nopython 模式。它在对象模式下工作，但这根本不会加速代码，所以它是我需要的 nopython 模式。

下面的代码只是一个说明问题的例子。我的实际代码有同样的问题，但有更多的循环和迭代，因此 jit 会非常有用。

import numpy as np
from numba import jit

n = 100
myarray = np.zeros(n)

@jit(nopython=True)
def compute(n):
    for i in xrange(n):
        myarray[i] += 1  # This indexing causes the error.

compute(n)

"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.SetItemConstrain object at 0x7f700c89a7d0>:
> Immutable array
> File "jit_test.py", line 10
"""

如果我尝试切片而不是出现不同的错误。

# The rest of the code unchanged.
        myarray[:] += 1  # This slicing causes the error.
# The rest of the code unchanged.
"""Sample run:
> python jit_test.py
> ...
> Failed at nopython (nopython frontend)
> Internal error at <numba.typeinfer.IntrinsicCallConstrain object at 0x7f297e2a9dd0>:
> Don't know how to create implicit output array with 'A' layout.
> File "jit_test2.py", line 10
"""

我的构建:

numba version: 0.19.1
numpy version: 1.9.2
python version: 2.7.10

Answer 1

在 numba 中，全局数组是静态的。您将在您的方法中获得一个只读副本。见:

http://numba.pydata.org/numba-doc/0.20.0/reference/pysemantics.html?highlight=global#global-and-closure-variables

如果您想修改一个数组，要么显式传入它，要么在函数中创建它(在进行数组内存管理的较新版本的 Numba 中)并返回它。

例如:

import numpy as np
import numba as nb

def compute(arr):
    for i in xrange(arr.shape[0]):
        arr[i] += 1

n = 100
myarray = np.zeros(n)

jitcompute = nb.jit(nopython=True)(compute)

然后计时:

In [12]: %timeit compute(myarray)
10000 loops, best of 3: 25.7 µs per loop

In [13]: %timeit jitcompute(myarray)
The slowest run took 17.06 times longer than the fastest. This could mean that an intermediate result is being cached
1000000 loops, best of 3: 461 ns per loop

一些注意事项:

• 这显然是一个简单的示例，但它展示了如何在方法之外定义一个数组，然后让函数就地修改它。
• 我没有在这里使用 @jit 装饰器，这样我就可以编写一次方法定义，并在使用和不使用 numba jit-ing 的情况下为方法计时，但是 jitcompute等同于在 compute 上使用装饰器会得到的结果。
• 如果您曾经对 numba 代码进行基准测试，请确保在计时之前运行该方法一次，否则您将看到的是 jit 代码所花费的时间和实际执行时间的组合。第二次运行时，您只会看到执行时间。