Numba 并行代码比顺序代码慢

Question

我是 Numba 的新手，我正在尝试使用 Numba（版本 0.54.1）在 Python 中实现旧的 Fortran 代码，但是当我添加 parallel = True 时，程序实际上变慢了吃下。我的程序非常简单：我在 L x L 网格中更改位置 x 和 y，然后对网格中的每个位置执行求和

import numpy as np
import numba as nb

@nb.njit(parallel=True)
def lyapunov_grid(x_grid, y_grid, k, N):
    L = len(x_grid)
    lypnv = np.zeros((L, L))
    for ii in nb.prange(L):
        for jj in range(L):
            x = x_grid[ii]
            y = y_grid[jj]
            beta0 = 0
            sumT11 = 0

            for j in range(N):
                y = (y - k*np.sin(x)) % (2*np.pi)
                x = (x + y) % (2*np.pi)
                J = np.array([[1.0, -k*np.cos(x)], [1.0, 1.0 - k*np.cos(x)]])
                beta = np.arctan((-J[1,0]*np.cos(beta0) + J[1,1]*np.sin(beta0))/(J[0,0]*np.cos(beta0) - J[0,1]*np.sin(beta0)))
                T11 = np.cos(beta0)*(J[0,0]*np.cos(beta) - J[1,0]*np.sin(beta)) - np.sin(beta0)*(J[0,1]*np.cos(beta) - J[1,1]*np.sin(beta))
                sumT11 += np.log(abs(T11))/np.log(2)

                beta0 = beta
            
            lypnv[ii, jj] = sumT11/N
    return lypnv

# Compile
_ = lyapunov_grid(np.linspace(0, 1, 10), np.linspace(0, 1, 10), 1, 10)
# Parameters
N = int(1e3)
L = 128
pi = np.pi
k = 1.5
# Limits of the phase space
x0 = -pi
xf = pi
y0 = -pi
yf = pi
# Grid positions
x = np.linspace(x0, xf, L, endpoint=True)
y = np.linspace(y0, yf, L, endpoint=True)

lypnv = lyapunov_grid(x, y, k, N)

parallel=False 大约需要 8 秒到运行，但是 parallel=True 大约需要 14 秒。我还使用 https://github.com/animator/mandelbrot-numba 中的另一个代码进行了测试，在这种情况下并行化有效。

import math
import numpy as np
import numba as nb

WIDTH = 1000
MAX_ITER = 1000

@nb.njit(parallel=True)
def mandelbrot(width, max_iter):     
    pixels = np.zeros((width, width, 3), dtype=np.uint8)
    for y in nb.prange(width):
        for x in range(width):
            c0 = complex(3.0*x/width - 2, 3.0*y/width - 1.5) 
            c = 0
            for i in range(1, max_iter): 
                if abs(c) > 2: 
                    log_iter = math.log(i) 
                    pixels[y, x, :] = np.array([int(255*(1+math.cos(3.32*log_iter))/2), 
                                                int(255*(1+math.cos(0.774*log_iter))/2), 
                                                int(255*(1+math.cos(0.412*log_iter))/2)], 
                                               dtype=np.uint8) 
                    break
                c = c * c + c0
    return pixels

# compile
_ = mandelbrot(WIDTH, 10)

calcpixels = mandelbrot(WIDTH, MAX_ITER)

Answer 1

一个主要问题是第二个函数调用再次编译函数。实际上，提供的参数的类型发生了变化：在第一次调用中，第三个参数是一个整数（int 转换为 np.int_），而在第二次调用中，第三个参数（k）是一个浮点数（float 转换为 np.float64）。 Numba 为不同的参数类型重新编译函数，因为它们是从参数的类型推导出来的，并且它不知道你想为第三个参数使用 np.float64 类型（因为第一次编译函数是为了np.int_ 类型）。解决此问题的一个简单方法是将第一个调用更改为：

_ = lyapunov_grid(np.linspace(0, 1, 10), np.linspace(0, 1, 10), 1.0, 10)

但是，这不是解决问题的可靠方法。您可以为 Numba 指定参数类型，以便它在声明时编译函数。这也消除了人为调用函数（使用无用参数）的需要。

@nb.njit('float64[:,:](float64[::1], float64[::1], float64, float64)', parallel=True)

请注意，(J[0,0]*np.cos(beta0) - J[0,1]*np.sin(beta0)) 第一次为零，导致除以 0。

另一个主要问题来自循环中许多小数组的分配，导致标准分配器争用（有关详细信息，请参阅）。虽然 Numba 理论上可以对其进行优化（即用局部变量替换数组），但实际上并不能，这会导致巨大的减速和争用。希望在您的情况下，您不需要实际创建数组。最后，您只能在包含循环中创建它并在最内层循环中修改它。这是优化后的代码：

@nb.njit('float64[:,:](float64[::1], float64[::1], float64, float64)', parallel=True)
def lyapunov_grid(x_grid, y_grid, k, N):
    L = len(x_grid)
    lypnv = np.zeros((L, L))
    for ii in nb.prange(L):
        J = np.ones((2, 2), dtype=np.float64)

        for jj in range(L):
            x = x_grid[ii]
            y = y_grid[jj]
            beta0 = 0
            sumT11 = 0

            for j in range(N):
                y = (y - k*np.sin(x)) % (2*np.pi)
                x = (x + y) % (2*np.pi)
                J[0, 1] = -k*np.cos(x)
                J[1, 1] = 1.0 - k*np.cos(x)
                beta = np.arctan((-J[1,0]*np.cos(beta0) + J[1,1]*np.sin(beta0))/(J[0,0]*np.cos(beta0) - J[0,1]*np.sin(beta0)))
                T11 = np.cos(beta0)*(J[0,0]*np.cos(beta) - J[1,0]*np.sin(beta)) - np.sin(beta0)*(J[0,1]*np.cos(beta) - J[1,1]*np.sin(beta))
                sumT11 += np.log(abs(T11))/np.log(2)

                beta0 = beta
            
            lypnv[ii, jj] = sumT11/N
    return lypnv

这是在旧的 2 核机器（有 4 个硬件线程）上的结果：

Original sequential:   15.9 s
Original parallel:     11.9 s
Fix-build sequential:  15.7 s
Fix-build parallel:    10.1 s
Optimized sequential:  2.73 s
Optimized parallel:    0.94 s

优化后的实现比其他实现快得多。与原始版本相比，并行优化版本的规模非常好（比顺序版本快 2.9 倍）。最后，最佳版本比原始并行版本快 12 倍。我希望在具有更多内核的最新机器上进行更快的计算。

Numba 并行代码比顺序代码慢

Numba parallel code slower than its sequential counterpart

python

parallel-processing

performance

numba