时间戳

首页 > 解决方案 > 如何使用 Python/Numba 对数组执行并行归约

python - 如何使用 Python/Numba 对数组执行并行归约

问题描述

我正在有效地寻找解决这个问题的方法。我了解将较大的数组减少为较小的数组（例如，如果您进行了某种粒子模拟，并且您正在将矩与网格点数组相加）会导致竞争条件，因为目标数组在被访问时会被线程锁定，导致其他线程丢弃值并导致错误输出。一种简单的解决方法是让每个线程在目标数组的副本上累积其值（因为可能只有数百个网格点，但有数百万个粒子），然后将这些数组添加到并行循环之外。

但是，我不知道从哪里开始。我的问题基本上归结为“如何为每个线程创建一个数组的本地副本，执行一些操作，并在最后总结它们？”。一个简化的例子可能是：

def deposit_particles_to_grid(ploc, pweight, grid):
    '''
    ploc    : (N)   : Particle location as an index in the grid array, leftmost point
    pweight : (Nx3) : Particle weights as a float across 3 gridpoints
    grid    : (M)   : Target array gridpoints where M << N
    '''
    for ii in nb.prange(ploc.shape[0]):
        grid[ploc[ii]]     += pweight[ii, 0]
        grid[ploc[ii] + 1] += pweight[ii, 1]
        grid[ploc[ii] + 2] += pweight[ii, 2]
    return

在这里，ploc.shape[0] 可能在几千万或几亿，而 grid 可能只有几百或几千个点。理想情况下，最好将 ploc 和 pweight 数组“分块”到 n_threads 中，并为每个线程提供自己的“网格”副本。我只是不确定该怎么做。

编辑示例以更能代表我当前的非并行代码。每个点不只是一对一的直方图，而是分布在 pweight 中定义的 3 个连续点上。

标签： pythonparallel-processingnumba

解决方案

回答我自己的问题，因为对于像我这样不熟练的人来说，这可能很有用。基本上我所做的只是为每个线程创建一个带有额外轴的新数组，并在最后在该轴上求和。这使值的沉积并行化（这是最密集的部分），并防止竞争条件，因为每个线程都有自己的轴。

我没有声明这是否可以大规模扩展或任何东西，但就我的目的而言，它工作得很好。不过，我欢迎对此进行任何改进！我知道如果预先分配了 grid_threads 内存（在可能多次调用此函数的情况下），则可以加快速度。

@nb.njit(parallel=True)
def deposit_particles_to_grid_parallel(ploc, pweight, grid):
    '''
    ploc    : (N)   : Particle location as an index in the grid array, leftmost point
    pweight : (Nx3) : Particle weights as a float across 3 gridpoints
    grid    : (M)   : Target array gridpoints where M << N
    '''
    grid_threads = np.zeros((grid.shape[0], n_threads), dtype=grid.dtype)
    N_per_thread = ploc.shape[0] / n_threads     
    n_start_idxs = np.arange(n_threads)*N_per_thread

    for tt in nb.prange(n_threads):
        for ii in range(n_start_idxs[tt], n_start_idxs[tt]+N_per_thread):
            grid_threads[ploc[ii],     tt] += pweight[ii, 0]
            grid_threads[ploc[ii] + 1, tt] += pweight[ii, 1]
            grid_threads[ploc[ii] + 2, tt] += pweight[ii, 2]
    grid[:] = grid_threads.sum(axis=1)
    return

推荐阅读