python - 有什么方法可以报告/共享 Numba prange 循环的结果吗？

Question

我有一个大的 prange 循环，可以在这里高效地执行一些过程。我想跟踪其中一个何时找到更好的结果。但是，似乎在 prange 循环内写入变量会创建一个新的私有(private)变量，因此无法以这种方式共享数据。

这是一个典型的例子:这个 prange 在多个数组中搜索最大的随机值。我希望他们“合作”，以便只打印新的最高值:

import numpy as np
from numba import jit, prange

@jit(nopython=True, parallel=True)
def test():
    arr = np.random.rand(1000, 15000000)

    largest_value = 0
    for i in prange(arr.shape[0]):
        for val in arr[i]:
            if val > largest_value:
                print('NEW LARGEST VALUE: ' + str(int(largest_value * 100000000)))
                largest_value = val

test()

运行它会打印伪值，例如:

...
NEW LARGEST VALUE: 85809181
NEW LARGEST VALUE: 93504938
NEW LARGEST VALUE: 75796671
NEW LARGEST VALUE: 83944676
NEW LARGEST VALUE: 99068766

这是因为每个线程在写入 largest_value 时都会创建一个本地私有(private)副本，因此他们无法以这种方式分享他们的发现。

是否有任何技巧可以使这成为可能？共享内存？有什么我想念的吗？还有别的吗？

Answer 1

线程之间共享数据的解决方案是在并行循环之前创建一个数组，这样线程就可以在并行循环中填充共享数组。请注意，您应该注意 false-sharing在共享数组中写入时(为了性能)。

在某些情况下，您只需要编写每个线程的信息。在这种情况下，您可以分配一个大小为 get_num_threads() 的数组(可能更大以防止错误共享)并在 get_thread_id() 提供的位置写入。这些函数应该由 Numba 在并行循环中提供。您可以在 documentation 中找到更多相关信息。 .

请注意，旧的/手动的方法是使用额外的包含循环 for threadId in prange(get_num_threads()) 并使用 threadId 作为索引共享数组。因此，必须手动划分迭代空间，以便在线程之间平均分担工作(这做起来有点麻烦)。幸运的是，不再需要此解决方案(至少在最简单的情况下)。