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转换为 a np.int_），而在第二次调用中，第三个参数 ( k) 是一个浮点数（float转换为 a 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 倍。我期望在具有更多内核的最新机器上计算速度更快。