如何使用 Jit (Numba) 并行化？

Question

我想与你分享我是如何解决我在 jit 中遇到的并行化问题的。最初我有下面的代码并得到了这个错误：

代码：

@jit(nopython=True, parallel=True)
def mandelbrot_2(cArray, iterations):
    count = 0
    for c in cArray:
        z = 0
        for n in range(iterations):
            z = np.square(z) + c
            if np.abs(z) > 2 : 
                count += 1 ; 
                break
    return cArray.shape[0] - count

错误：

/Users/alexander/opt/anaconda3/lib/python3.8/site-packages/numba/core/typed_passes.py:326: NumbaPerformanceWarning: [1m
The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.

To find out why, try turning on parallel diagnostics, see https://numba.pydata.org/numba-doc/latest/user/parallel.html#diagnostics for help.
[1m
File "<ipython-input-56-9fcfc0a9fe03>", line 17:[0m
[1m@jit(nopython=True, parallel=True)
[1mdef mandelbrot_2(cArray, iterations):
[0m[1m^[0m[0m
[0m
  warnings.warn(errors.NumbaPerformanceWarning(msg,

Answer 1

1. 导入范围： from numba import jit, prange

2. 新代码：

• 将第一个循环转换为索引数组，而不是每次都使用新的数组元素。
• 将两个循环中的 range 转换为我从 Numba 导入的 prange。

@jit(nopython=True, parallel=True)
def mandelbrot_2(cArray, iterations):
    count = 0
    for c in prange(cArray.shape[0]):
        z = 0
        for n in prange(iterations):
            z = np.square(z) + cArray[c]
            if np.abs(z) > 2 : count += 1 ; break
    return cArray.shape[0] - count