python - 如何使用 Numba 并行化数组上的 for 循环

Question

我正在尝试编写一个 Numba 函数，该函数使用 @numba.njit(parallel=True) 循环遍历列表(而不是范围)。 。例如，

import numpy as np
import numba

arr = np.ones(10)
idx = np.array([4, 2, 5])

@numba.njit(parallel=True)
def foo(arr, idx):
    for i in idx:
        arr[i] = 0
    return arr

foo(arr, idx)

我收到的警告是NumbaPerformanceWarning: The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.我知道可以使用 numba.prange() 显式并行化类似的循环。 ，但我需要循环来遍历索引数组。这可能吗？如何实现？

Answer 1

for i in idx: ... 在语义上相当于:

for j in range(len(idx)):
    i = idx[j]
    [...]

第二个版本的性能可能略有不同。在第二个版本的基础上，您可以将 range 替换为 prange 并获得并行版本。

但是，请注意，如果 idx 的两个值相同，这将导致竞争条件。竞争条件会导致未定义的行为(基本上，任何事情都可能发生)。