如何并行化Numba实现的矩阵负元素索引遍历循环？

问题背景 #

需要从大型矩阵中提取所有小于0的元素的索引，已实现单线程Numba代码，比原生np.where快，但略慢于用@nb.njit包装的np.where，希望通过并行化进一步优化，核心难点是处理循环中共享的idx索引避免竞争。

原代码如下：

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]

并行化优化方案 #

直接用nb.prange遍历会导致多个线程同时修改idx，引发数据竞争和错误结果。正确的做法是分两步并行：先统计每个行的负元素数量，再通过前缀和确定每个线程的写入区间，最后并行写入索引。

优化后的代码：

import numba as nb
import numpy as np

A = np.random.random(size=(1000, 1000)) - 0.1

@nb.njit(parallel=True, cache=True)
def numba_parallel(arr):
    rows_total = arr.shape[0]
    cols_total = arr.shape[1]
    # 第一步：统计每行的负元素数量
    count_per_row = np.zeros(rows_total, dtype=np.int64)
    for i in nb.prange(rows_total):
        cnt = 0
        for j in range(cols_total):
            if arr[i, j] < 0:
                cnt += 1
        count_per_row[i] = cnt
    
    # 计算前缀和，确定每行的起始写入位置
    prefix_sum = np.zeros(rows_total + 1, dtype=np.int64)
    for i in range(rows_total):
        prefix_sum[i+1] = prefix_sum[i] + count_per_row[i]
    
    total_neg = prefix_sum[-1]
    rows = np.empty(total_neg, dtype=np.int64)
    cols = np.empty(total_neg, dtype=np.int64)
    
    # 第二步：并行写入索引
    for i in nb.prange(rows_total):
        start_idx = prefix_sum[i]
        current_idx = start_idx
        for j in range(cols_total):
            if arr[i, j] < 0:
                rows[current_idx] = i
                cols[current_idx] = j
                current_idx += 1
    
    return rows, cols

关键优化点 #

• 避免共享变量竞争：通过统计每行负元素数量+前缀和，让每个线程只负责自己行的索引写入，无需修改全局共享的idx
• 并行模式开启：添加parallel=True参数，配合nb.prange实现循环并行
• 内存预分配优化：根据统计的总负元素数量分配数组，避免原代码中过度预分配的内存浪费
• 数据类型优化：索引用np.int64而非默认浮点型，减少内存占用和转换开销

性能说明 #

该并行版本在多核CPU上通常能比原单线程Numba代码快2-8倍（取决于CPU核心数），性能接近甚至超过@nb.njit包装的np.where，尤其在超大矩阵场景下优势更明显。

内容的提问来源于stack exchange，提问作者gaussplustwo