Numba并行装饰器下含竞争条件的例程问题求解

你编写了一个并行计算不同随机矩阵SVD的Numba例程，代码如下：

import numpy as np 
from numba import jit,prange 

@jit(nopython=True,parallel=True)
def svd_bn(aa,n):
    
    res=[]
    for k in prange(n):
        u,s,v=np.linalg.svd(aa[k],0)
        res.append( s ) 
        
        
    return res
aa=[np.random.rand(2*k,2*(k+1)) for k in range(1,10)]

res=svd_bn(aa,len(aa))

该代码在Windows系统可正常运行并返回有序列表，但在Linux系统出现double free or corruption (!prev)错误。问题根源是Numba并行环境中Python列表的append操作非线程安全，动态修改列表会触发内存管理竞争，Linux的内存检测机制更严格，因此暴露了这个问题。以下是针对无法预先知晓数组形状的解决方案：

方法一：使用Numba类型化列表（推荐） #

Numba的typed.List是线程安全的类型化容器，适合存储形状不确定的数组，且能在nopython模式下高效运行：

import numpy as np 
from numba import jit, prange
from numba.typed import List

@jit(nopython=True, parallel=True)
def svd_bn(aa, n):
    # 初始化类型化列表，预分配n个空元素避免并行扩容竞争
    res = List()
    for _ in range(n):
        res.append(np.empty(0, dtype=np.float64))
    
    for k in prange(n):
        u, s, v = np.linalg.svd(aa[k], 0)
        # 通过索引直接赋值，替代线程不安全的append
        res[k] = s
    
    return res

# 将输入列表转换为Numba类型化列表，适配nopython模式
aa = List()
for k in range(1, 10):
    aa.append(np.random.rand(2*k, 2*(k+1)))

res = svd_bn(aa, len(aa))

方法二：用对象数组暂存结果（兼容旧版Numba） #

如果无法使用typed.List，可以先通过对象数组在并行循环中存储每个结果，最后单线程转换为列表：

import numpy as np 
from numba import jit, prange
from numba.typed import List

@jit(nopython=True, parallel=True)
def svd_bn(aa, n):
    # 预分配对象数组，用于存储不同形状的s数组
    res_arr = np.empty(n, dtype=object)
    
    for k in prange(n):
        u, s, v = np.linalg.svd(aa[k], 0)
        res_arr[k] = s
    
    # 单线程环境下转换为列表，避免并行修改竞争
    res = []
    for k in range(n):
        res.append(res_arr[k])
    return res

# 转换输入为类型化列表
aa = [np.random.rand(2*k, 2*(k+1)) for k in range(1, 10)]
aa_typed = List()
for arr in aa:
    aa_typed.append(arr)

res = svd_bn(aa_typed, len(aa_typed))

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