如何并行化矩阵上的简单循环？

Question

我有一个大矩阵，我想输出矩阵中元素小于 0 的所有索引。我在 numba 中有这个 MWE：

import numba as nb
import numpy as np
A = np.random.random(size = (1000, 1000)) - 0.1
@nb.njit(cache=True)
def numba_only(arr):
    rows = np.empty(arr.shape[0]*arr.shape[1])
    cols = np.empty(arr.shape[0]*arr.shape[1])
    idx = 0
    for i in range(arr.shape[0]):
        for j in range(A.shape[1]):
            if arr[i, j] < 0:
                rows[idx] = i
                cols[idx] = j
                idx += 1
    return rows[:idx], cols[:idx]

时间，我得到：

%timeit numba_only(A)
2.29 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这比

np.where(A<0)

（没有 numba）更快，它给出：

%timeit numpy_only(A)
3.56 ms ± 59.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

但是比被

np.where

包裹的

@nb.njit

慢一点：

@nb.njit(cache=True)
def numba_where(arr):
    return np.where(arr<0)

给出：

%timeit numba_where(A)
1.76 ms ± 84.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba 代码可以通过并行化来加速吗？我意识到它可能受内存限制，但我认为现代硬件应该允许某种程度的并行访问内存。

由于循环中的索引 idx，我不确定如何使用 nb.prange 来实现此目的。

Answer 1

由于指令的长依赖链也依赖于数据，所以所提供的算法本质上是顺序的。但是，可以使用 scan 并行执行等效操作。

实现这种基于扫描的并行算法要复杂得多（尤其是高效的缓存友好扫描）。一个相对简单的实现是在数据集上迭代两次，以便计算局部和，然后计算累积和，最后完成实际工作。

@nb.njit(cache=True, parallel=True)
def numba_parallel_where(arr):
    flattenArr = arr.reshape(-1)
    arrSize = flattenArr.size
    chunkSize = 2048
    chunkCount = (arrSize + chunkSize - 1) // chunkSize
    counts = np.empty(chunkCount, dtype=np.int32)
    for chunkId in nb.prange(chunkCount):
        start = chunkId * chunkSize
        end = min(start + chunkSize, arrSize)
        count = 0
        for i in range(start, end):
            count += flattenArr[i] < 0
        counts[chunkId] = count
    offsets = np.empty(chunkCount + 1, dtype=np.int64)
    offsets[0] = 0
    for chunkId in range(chunkCount):
        offsets[chunkId + 1] = offsets[chunkId] + counts[chunkId]
    outSize = offsets[-1]
    rows = np.empty(outSize, dtype=np.int32)
    cols = np.empty(outSize, dtype=np.int32)
    n = np.int32(arr.shape[1])
    for chunkId in nb.prange(chunkCount):
        start = chunkId * chunkSize
        end = min(start + chunkSize, arrSize)
        idx = offsets[chunkId]
        for i in range(start, end):
            if flattenArr[i] < 0:
                rows[idx] = i // n
                cols[idx] = i % n
                idx += 1
    return rows, cols

---

结果

以下是我的 6 核 i5-9600KF CPU 和 40 GiB/s RAM（最佳读取带宽）在 Windows 上的性能结果：

Sequential:              1.68 ms
Parallel:                0.85 ms   <-----
Theoretical optimal:     0.20 ms

并行实现大约快两倍。在我的机器上，它成功达到了 21 GiB/s 的内存吞吐量，虽然不是最优的，但也不算太糟糕。虽然这在 6 核 CPU 上看起来令人失望，但并行实现并不是非常优化，有几点需要考虑：

• 分配输出真的很慢（似乎需要0.20-0.25毫秒）
• 输入数组被读取两次给内存带来很大压力
• 模/除法用于计算索引有点昂贵

---

注释

可能的优化包括：

• 预分配输出（手动填充零以避免可能的页面错误）；
• 使用块集群以避免在整个数据集上迭代两次（这在 Numba 中很难编写快速实现，因为与本机语言相比，并行性非常有限）；
• 迭代矩阵的行以避免取模/除法，假设行数出于性能考虑既不太大也不太小。

请注意，此类优化将使实现比上述优化更加复杂。