Numba 并行代码比顺序代码慢

我是 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)

With 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)

一个主要问题是第二个函数调用再次编译该函数。事实上，提供的参数的类型发生了变化：在第一次调用中，第三个参数是一个整数（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。

另一个主要问题来自于循环中许多小数组的分配导致标准分配器的争用 (see 这个帖子了解更多信息）。虽然 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 倍比原来的并行版本。我期望在具有更多内核的最新机器上计算速度更快。