如何用Numba优化随机Runge-Kutta？并行提速不达预期问题排查

我正在优化一个随机微分方程求解任务，需要独立求解Mmax个微分方程。用Numba加速后，Mmax=10000时耗时仅从16秒降到14秒，远低于预期优化效果，代码无报错但明显存在性能瓶颈，求分析原因。

核心求解代码 #

@njit(fastmath=True)
def Liouvillian(rho,t_list,alpha1,alpha1p):
    C = np.zeros((len(rho0)),dtype=np.complex64)
    B = L0+alpha1p*L1m+alpha1*L1p
    for j in range(len(rho0)):
        for i in range(len(rho0)):
            C[j] += B[j, i]*rho[i] 
    return C

@njit(parallel=True, fastmath=True)
def RK4(rho0,t_list,alpha1,alpha1p):
    rho=np.zeros((len(rho0),len(t_list),Mmax),dtype=np.complex64) 
    for m in prange(Mmax):
        rho[:,0,m]=rho0
        n=0
        for t in t_list[0:len(t_list)-1]:
            k1=Liouvillian(rho[:,n,m],t,alpha1[st*n,m],alpha1p[st*n,m])
            k2=Liouvillian(rho[:,n,m]+dt*k1/2,t+dt/2,alpha1[st*n,m],alpha1p[st*n,m])
            k3=Liouvillian(rho[:,n,m]+dt*k2/2,t+dt/2,alpha1[st*n,m],alpha1p[st*n,m])
            k4=Liouvillian(rho[:,n,m]+dt*k3,t+dt,alpha1[st*n,m],alpha1p[st*n,m])
            rho[:,n+1,m]=rho[:,n,m]+dt*(k1+2*k2+2*k3+k4)/6.0
            n=n+1
    return rho

说明：alpha1p和alpha1是[t,m]维度的二维数组，L0、L1m、L1p为预定义普通数组。

辅助参数生成代码 #

def A_matrix(num_modes,kappa,E_0):
    A_modes=np.zeros(((2*num_modes),(2*num_modes)),dtype=np.complex64)
    for i in range(np.shape(A_modes)[0]):
        for j in range(np.shape(A_modes)[0]):
            if i==j:
                A_modes[i,j]=-kappa/4
                if i==0:
                    A_modes[i,i+num_modes+1]=E_0*kappa/2  
                    A_modes[i,-1]=E_0*kappa/2 
                elif 0<i<num_modes-1:
                    A_modes[i,i+num_modes+1]=E_0*kappa/2
                    A_modes[i,i+num_modes-1]=E_0*kappa/2 
                elif i==num_modes:
                    A_modes[i-1,i]=E_0*kappa/2 
                    A_modes[i-1,-2]=E_0*kappa/2 
    return (A_modes+A_modes.transpose())


def D_mat(num_modes):
    D_matrix=np.zeros(((num_modes),(num_modes)),dtype=np.complex64)
    for i in range(np.shape(D_matrix)[0]):
        for j in range(np.shape(D_matrix)[0]):
            if i==j:
                if i==0:
                    D_matrix[i,i+1]=1
                    D_matrix[i,-1]=1
                elif 0<i<num_modes-1:
                    D_matrix[i,i+1]=1
    return (D_matrix+D_matrix.transpose())

def a_part(alpha,t_list):
    M=A_matrix(num_modes,kappa,E_0)
    return M.dot(alpha)


def b_part(w,t_list):
    D_1=D_mat(num_modes)
    Zero_modes=np.zeros(((num_modes),(num_modes)),dtype=np.complex64)
    D=np.bmat([[D_1, Zero_modes], [Zero_modes, D_1]])
    B=sqrtm(D)*np.sqrt(E_0*kappa/2)
    return B.dot(w)

def SDE_Param_Euler_Mauyrama(Mmax):
    alpha=np.zeros((2*num_modes,Smax+1,Mmax),dtype=np.complex64)
    n=0
    alpha[:,n,:]=0.0+1j*0.0
    for s in s_list[0:len(s_list)-1]:
        alpha[:,n+1,:]=alpha[:,n,:]+ds*a_part(alpha[:,n,:],s)+b_part(w[:,n,:],s)
        n=n+1
    return (alpha)

#Parameters
E_0=0.5
kappa=10.
gamma=1.
num_modes=2 ## number of modes

Mmax=10000 #number of samples

Tmax=20 ##max value for time
dt=1/(2*kappa)

st=10

ds=dt/st
Nmax=int(Tmax/dt) ##number of steps
Smax=int(Tmax/ds) ##number of steps

t_list=np.arange(0,Tmax+dt/2,dt)
s_list=np.arange(0,Tmax+ds/2,ds)


w = np.random.randn(2*num_modes,Smax+1,Mmax)*np.sqrt(ds)

(alpha1,alpha2,alpha1p,alpha2p)=SDE_Param_Euler_Mauyrama(Mmax)

from qutip import *
Delta_a=0

num_qubits=1
##Atom 1
sm_1=sigmam()
sp_1=sigmap()
sx_1=sigmax()
sy_1=sigmay()
sz_1=sigmaz()


state0=np.zeros(2**num_qubits,dtype=np.complex64)
state0[-1]=1

rho0=np.kron(state0,np.conjugate(np.transpose(state0)))

Ha=Delta_a*(sz_1)/2    #Deterministic Liouvillian
L_ops=[np.sqrt(gamma)*sm_1]

L0=np.array(liouvillian(Ha,L_ops),dtype=np.complex64)

L1m=-np.array(np.sqrt(kappa*gamma)*1j*liouvillian(sm_1),dtype=np.complex64)
L1p=np.array(np.sqrt(kappa*gamma)*1j*liouvillian(sp_1),dtype=np.complex64)

rho_1=RK4(rho0,t_list,alpha1,alpha1p).reshape(2**num_qubits,2**num_qubits,len(t_list),Mmax)

问题分析与优化建议 #

1. 核心性能瓶颈：手动矩阵乘法效率极低 #

Liouvillian中用双重循环手动实现矩阵向量乘法，这是最大的性能浪费。Numba虽然能加速循环，但远不如NumPy内置的BLAS优化矩阵乘法高效，直接替换成B @ rho或np.dot(B, rho)，性能会有质的提升。

2. 并行化的内存访问问题 #

• rho数组维度为(len(rho0), len(t_list), Mmax)，在prange(Mmax)循环中，线程访问rho[:,n,m]属于非连续跨步访问（NumPy默认列主序），缓存命中率极低，严重拖慢并行效率。
• 建议调整数组维度为(Mmax, len(t_list), len(rho0))，让每个线程连续访问自身样本数据，充分利用CPU缓存。

3. 全局变量依赖拖慢编译与执行 #

Liouvillian直接使用全局变量rho0、L0、L1m、L1p，Numba编译时需额外处理全局变量访问，增加编译时间且可能导致运行时间接访问开销，应将这些变量作为参数传入函数。

4. 不必要的参数传递 #

Liouvillian的t_list参数未被使用，可直接删除，减少参数传递开销。

优化后的核心代码示例 #

@njit(fastmath=True)
def Liouvillian(rho, alpha1, alpha1p, L0, L1m, L1p):
    B = L0 + alpha1p * L1m + alpha1 * L1p
    return B @ rho  # 直接用矩阵向量乘法

@njit(parallel=True, fastmath=True)
def RK4(rho0, t_list, alpha1, alpha1p, L0, L1m, L1p, dt, st, Mmax):
    # 调整数组维度，优化内存访问
    rho = np.zeros((Mmax, len(t_list), len(rho0)), dtype=np.complex64)
    for m in prange(Mmax):
        rho[m, 0, :] = rho0
        for n in range(len(t_list)-1):
            current_alpha1 = alpha1[st*n, m]
            current_alpha1p = alpha1p[st*n, m]
            k1 = Liouvillian(rho[m, n, :], current_alpha1, current_alpha1p, L0, L1m, L1p)
            k2 = Liouvillian(rho[m, n, :] + dt*k1/2, current_alpha1, current_alpha1p, L0, L1m, L1p)
            k3 = Liouvillian(rho[m, n, :] + dt*k2/2, current_alpha1, current_alpha1p, L0, L1m, L1p)
            k4 = Liouvillian(rho[m, n, :] + dt*k3, current_alpha1, current_alpha1p, L0, L1m, L1p)
            rho[m, n+1, :] = rho[m, n, :] + dt*(k1 + 2*k2 + 2*k3 + k4)/6.0
    # 最后调整回原维度
    return rho.transpose(2, 1, 0)

5. 其他优化点 #

• 预计算alpha1[st*n, m]和alpha1p[st*n, m]，避免重复索引计算。
• SDE_Param_Euler_Mauyrama中的循环同样可用Numba加速，这部分也是耗时操作。
• 用np.ascontiguousarray确保数组内存连续，进一步提升缓存效率。

按上述优化点修改后，并行效率会显著提升，耗时可降到几秒甚至更低。

内容的提问来源于stack exchange，提问作者J.Agusti