启用parallel=True时Numba致Python崩溃，求高性能并行方案

问题描述 #

尝试用Numba加速N体问题的Leapfrog积分计算，给nnleapfrog_integrate函数的jit装饰器添加parallel=True标志时，调用函数会触发段错误导致Python进程崩溃。移除该标志后函数可正常运行，但仅以单线程模式执行，无法充分利用CPU核心，需要解决并行化后的段错误问题以获得最优性能。

出错代码 #

from numba import jit, prange
import numpy as np

@jit('Tuple((f8[:,:,::1],f8[:,:,::1]))(f8[:,::1], f8[:,::1], f8[::1], i8, i8, i8, f8, f8)', nopython=True,  parallel=True)
def nnleapfrog_integrate(pos, vel, mass, i_steps, r_steps, dt, G, softening):
    N = pos.shape[0]
    pos_data = np.zeros((int(np.ceil(i_steps/r_steps)), N, 3))
    vel_data = np.zeros((int(np.ceil(i_steps/r_steps)), N, 3))
    data_idx = 0
    acc = np.zeros((N,3))
    for s in range(i_steps):
        vel += acc * dt/2.0
        pos += vel * dt
        for i in prange(N):
            acc[i,0] =  0
            acc[i,1] =  0
            acc[i,2] =  0
            for j in range(N):
                dx = pos[j,0] - pos[i,0]
                dy = pos[j,1] - pos[i,1]
                dz = pos[j,2] - pos[i,2]
                inv_r3 = (dx**2 + dy**2 + dz**2 + softening**2)**(-1.5)
                acc[i,0] +=  G * (dx * inv_r3) * mass[j]
                acc[i,1] +=  G * (dy * inv_r3) * mass[j]
                acc[i,2] +=  G * (dz * inv_r3) * mass[j]
        vel += acc * dt/2.0
        if s % r_steps == 0:
            pos_data[data_idx] = pos
            vel_data[data_idx] = vel
            data_idx += 1
    return pos_data, vel_data

N = 10
dt = 60
pos = np.random.rand(N, 3)
vel = np.random.rand(N, 3)
m = np.random.rand(N)
softening = 1e3
G = 6.67430e-11
t_max = 3600*24*30
i_steps = int(t_max/dt)
r_steps = int(3600*24/dt)
r_i, v_i = nnleapfrog_integrate(pos, vel, m, i_steps, r_steps, dt, G, softening)

已尝试方案 #

将适合并行化的for i in prange(N):循环拆分到独立的getAcc函数中，该函数添加parallel=True后可正常运行并利用全部CPU核心，但整体性能比原函数的单线程版本慢3倍以上。

拆分后的代码：

from numba import jit, prange
import numpy as np

@jit('f8[:, ::1](f8[:, ::1], f8[::1], f8, f8)', nopython=True, parallel=True)
def getAcc( pos, mass, G, softening ):
    N = pos.shape[0]
    a = np.zeros((N,3))
    for i in prange(N):
        for j in range(N):
            dx = pos[j,0] - pos[i,0]
            dy = pos[j,1] - pos[i,1]
            dz = pos[j,2] - pos[i,2]
            inv_r3 = (dx**2 + dy**2 + dz**2 + softening**2)**(-1.5)
            a[i,0] +=  G * (dx * inv_r3) * mass[j]
            a[i,1] +=  G * (dy * inv_r3) * mass[j]
            a[i,2] +=  G * (dz * inv_r3) * mass[j]
    return a

@jit('Tuple((f8[:,:,::1],f8[:,:,::1]))(f8[:,::1], f8[:,::1], f8[::1], i8, i8, i8, f8, f8)', nopython=True)
def nleapfrog_integrate(pos, vel, mass, i_steps, r_steps, dt, G, softening):
    N = pos.shape[0]
    pos_data = np.zeros((int(np.ceil(i_steps/r_steps)), N, 3))
    vel_data = np.zeros((int(np.ceil(i_steps/r_steps)), N, 3))
    data_idx = 0
    acc = getAcc(pos, mass, G, softening)
    for i in range(i_steps):
        vel += acc * dt/2.0
        pos += vel * dt
        acc = getAcc( pos, mass, G, softening )
        vel += acc * dt/2.0
        if i % r_steps == 0:
            pos_data[data_idx] = pos
            vel_data[data_idx] = vel
            data_idx += 1
    return pos_data, vel_data

N = 10
dt = 60
pos = np.random.rand(N, 3)
vel = np.random.rand(N, 3)
m = np.random.rand(N)
softening = 1e3
G = 6.67430e-11
t_max = 3600*24*30
i_steps = int(t_max/dt)
r_steps = int(3600*24/dt)
r_i, v_i = nleapfrog_integrate(pos, vel, m, i_steps, r_steps, dt, G, softening)

性能测试结果 #

In [4]: %timeit r_i, v_i = nleapfrog_integrate(pos, vel, m, i_steps, r_steps, dt, G, softening)
8.51 s ± 46.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [5]: %timeit r_i, v_i = nnleapfrog_integrate(pos, vel, m, i_steps, r_steps, dt, G, softening)
2.53 s ± 18.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此，需要解决原函数内嵌prange循环并行化时的段错误问题，以实现最优性能。

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