python - 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 (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.

另一个主要问题来自循环中许多小数组的分配导致标准分配器的争用(有关详细信息，请参阅 this post)。虽然 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 倍。我希望在具有更多内核的最新机器上进行更快的计算。