20核Numba并行代码慢于串行的优化问题求助

我希望为科学计算函数get_Cilmoft添加@njit(parallel=True, fastmath=True)（尝试搭配nogil=True）实现并行加速，在20核环境下通过NUMBA_NUM_THREADS=20 python main.py运行，却发现并行版本执行速度显著慢于串行版本。

使用timeit.Timer对函数进行性能测试，测试代码如下：

if __name__ == "__main__":
    from timeit import Timer
    eval_legendre_withNJIT(2)
    get_Cilmoft(Psi)
    t = Timer(lambda: get_Cilmoft(Psi))
    print( min( t.repeat(repeat=2, number=2) ) )

@njit(parallel=True, fastmath=True)
def get_Cilmoft(Psi):
    Cilmoft = np.zeros( (gl.N_rs - 1,  gl.N_thetas,  2*gl.N_thetas-1,  2), dtype=np.complex128)
    for i in prange(Cilmoft.shape[0]):
        for l in prange(Cilmoft.shape[1]):
            for m in prange(-l, l+1, 1):
                exp_of_phis = np.exp(-1j * m * gl.phis)
                for contor in prange(Psi.shape[2]):
                    Psi[:, :, contor] = Psi[:, :, contor] * exp_of_phis[contor]
                F1 = np.zeros( (Psi.shape[0], Psi.shape[1]), dtype=np.complex128)
                for a in prange(Psi.shape[0]):
                    for b in prange(Psi.shape[1]):
                        F1[a, b] = np.trapz(Psi[a, b, :])
                poly_results = eval_legendre_withNJIT(l)
                F2 = np.zeros( (gl.N_rs - 1, ), dtype=np.complex128)
                for idx in prange(gl.N_rs - 1):
                    F2[idx] = np.sum(weights_here * poly_results * F1[idx, :])
                Cilmoft[i, l, m+l, 0] = np.complex128(np.sum(eigenvectors_memory[:, i, l] * F2 * lobatto_weights))
    return Cilmoft

• 将循环内临时数组F1、F2移至函数开头分配，无明显性能提升；
• 仅保留最外层循环使用prange，其余循环用range，性能略有提升但仍慢于串行；
• 添加nogil=True，执行时间变化极小；
• 替换np.sum为手动累加循环，无明显改善；
• 将循环迭代数改为常量，替代动态形状获取，无明显改善；
• 预计算exp_of_phis_matrix和poly_results_vector，性能首次提升但仍慢于串行；
• 自定义@njit(fastmath=True)修饰的sum和trapz函数替代原numpy函数，并行版本性能反超串行，不同N_r下的性能数据如下：
  • N_r=80：并行14.42s，串行22.23s
  • N_r=160：并行69.47s，串行104.85s
  • N_r=320：并行312.89s，串行513.18s
  • N_r=640：并行1365.2s，串行1648.88s

包含main.py预代码、gl.py配置文件、特征值/特征向量生成代码（与原代码一致）。

1. 初始Numba并行实现性能远低于串行的原因 #

• 过度嵌套并行：初始代码在多层循环（i、l、m、contor、a、b、idx）都使用prange，线程调度开销远大于并行计算收益。内层循环计算量小，线程频繁切换的成本直接抵消甚至超过加速效果。
• Numpy函数与并行环境冲突：循环内调用的np.trapz、np.sum是Python层实现，在Numba并行区域会触发额外GIL操作或同步开销，甚至强制串行执行，导致并行失效。
• 共享数据竞争：多线程同时修改全局数组Psi，引发缓存一致性问题，线程需等待同步，大幅降低执行效率。
• 临时数组频繁分配：循环内反复创建F1、F2等数组，并行环境下内存分配的竞争和回收开销被放大。

2. 进一步优化方向 #

• 精准控制并行粒度：仅在计算量最大的外层循环（如i循环）使用prange，避免嵌套并行，减少线程调度开销。
• 预计算静态数据：提前计算所有l对应的Legendre多项式结果并存储为数组，避免循环内调用eval_legendre_withNJIT；预计算所有m对应的exp_of_phis矩阵，直接索引使用。
• 内存布局优化：调整数组维度为C语言行优先（将频繁访问的维度前置），减少CPU缓存miss；提前分配所有临时数组并复用，避免循环内内存分配。
• 替换Numpy内置函数：继续用自定义Numba优化函数替代np.sum、np.trapz，甚至手动实现累加和积分逻辑，让Numba完全编译优化，避免Python层开销。
• 减少全局变量依赖：将gl.N_rs、gl.phis等全局变量作为参数传入函数，帮助Numba更好地做类型推断和优化，避免全局变量访问的额外开销。
• 利用SIMD指令：对内层乘法、累加操作，使用Numba的@vectorize装饰器或手动展开循环，充分利用CPU的SIMD向量指令，提升单线程计算效率。
• 消除数据依赖：检查循环间的数据依赖，比如用局部副本替代对Psi的共享修改，避免多线程间的竞争，进一步提升并行效率。

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