Numba 0.54.1开启parallel=True编译函数返回错误结果问题

经确认，下述问题是numba（0.54.1版本）的bug，目前Numba开发团队正在排查。

问题描述 #

我有一个简易的Python/Numpy函数，需要通过Numba编译加速。使用默认参数编译时运行完全正常，但将parallel=True传入njit装饰器后，返回无意义的损坏数据，暂未定位到问题根源。

可复现问题的示例代码 #

import numpy as np
from numba import njit, prange

@njit(boundscheck=False, nogil=True, parallel=False)
def ubcm_jacobian(
    D: np.ndarray,
    b: np.ndarray
) -> np.ndarray:
    """UBCM负对数似然的雅可比矩阵

    参数
    ----------
    D
        度序列数组
    b
        模型参数
    """
    B = np.empty_like(b)
    b = np.exp(b)

    for i in prange(len(b)):
        bi   = b[i]
        x    = bi * b
        B[i] = (x / (1 + x)).sum()

    b2 = b**2
    B -= b2 / (1 + b2)

    return B - D

测试用例 #

D = np.arange(1, 11)
b = np.log(D)

ubcm_jacobian(D, b).sum()   # 正确结果约为28.9899

异常表现 #

若将函数的编译参数改为parallel=True，返回结果完全异常，表现近似主循环未对B数组完成赋值。
该逻辑属于非常容易并行的场景，请问是什么原因导致该问题？

Numba并行诊断日志 #

================================================================================
 Parallel Accelerator Optimizing:  Function ubcm_jacobian, <ipython-
input-128-8dd321ee9df4> (1)  
================================================================================


Parallel loop listing for  Function ubcm_jacobian, <ipython-input-128-8dd321ee9df4> (1) 
----------------------------------------------------------|loop #ID
@njit(boundscheck=False, nogil=True, parallel=True)       | 
def ubcm_jacobian(                                        | 
    D: np.ndarray,                                        | 
    b: np.ndarray                                         | 
) -> np.ndarray:                                          | 
    """Jacobian of UBCM negative log-likelihood.          | 
                                                          | 
    Parameters                                            | 
    ----------                                            | 
    D                                                     | 
        Degree sequence array.                            | 
    b                                                     | 
        Model parameters.                                 | 
    """                                                   | 
    B = np.empty_like(b)                                  | 
    b = np.exp(b)-----------------------------------------| #39
                                                          | 
    for i in prange(len(b)):------------------------------| #46
        bi   = b[i]                                       | 
        x    = bi * b-------------------------------------| #40
        B[i] = (x / (1 + x)).sum()------------------------| #41, 45
                                                          | 
    b2 = b**2---------------------------------------------| #42
    B -= b2 / (1 + b2)------------------------------------| #43, 47
                                                          | 
    return B - D------------------------------------------| #44
------------------------------ After Optimisation ------------------------------
Parallel region 0:
+--39 (parallel)
   +--40 (serial, fused with loop(s): 41, 45)


 
Parallel region 0 (loop #39) had 2 loop(s) fused and 1 loop(s) serialized as 
part of the larger parallel loop (#39).
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

升级到numba最新版本0.54.1后结果有所变化，但仍不正确：开启parallel=True时返回结果约为31.99498，和parallel=False得到的正确结果28.9899仍有偏差，问题依旧存在。

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