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代码所花费的时间与实际执行时间的组合.第二次运行它,你只会看到执行时间.