numba - guvectorize 只比 jit 快一点

我试图并行化在许多独立数据集上运行的蒙特卡罗模拟。我发现 numba 的并行 guvectorize 实现仅比 numba jit 实现快 30-40%。

我找到了这些（1 https://stackoverflow.com/questions/31101699/numba-guvectorize-is-very-slow-compared-to-jit, 2 https://stackoverflow.com/questions/35350689/numba-guvectorize-target-parallel-slower-than-target-cpu）Stackoverflow 上的类似主题，但它们并没有真正回答我的问题。在第一种情况下，由于回退到对象模式，实现速度变慢，而在第二种情况下，原始发布者没有正确使用 guvectorize - 这些问题都不适用于我的代码。

为了确保我的代码没有问题，我创建了这段非常简单的代码来比较 jit 和 guvectorize：

import timeit
import numpy as np
from numba import jit, guvectorize

#both functions take an (m x n) array as input, compute the row sum, and return the row sums in a (m x 1) array

@guvectorize(["void(float64[:], float64[:])"], "(n) -> ()", target="parallel", nopython=True)
def row_sum_gu(input, output) :
    output[0] = np.sum(input)

@jit(nopython=True)
def row_sum_jit(input_array, output_array) :
    m, n = input_array.shape
    for i in range(m) :
        output_array[i] = np.sum(input_array[i,:])

rows = int(64) #broadcasting (= supposed parallellization) dimension for guvectorize
columns = int(1e6)
input_array = np.ones((rows, columns))
output_array = np.zeros((rows))
output_array2 = np.zeros((rows))

#the first run includes the compile time
row_sum_jit(input_array, output_array)
row_sum_gu(input_array, output_array2)

#run each function 100 times and record the time
print("jit time:", timeit.timeit("row_sum_jit(input_array, output_array)", "from __main__ import row_sum_jit, input_array, output_array", number=100))
print("guvectorize time:", timeit.timeit("row_sum_gu(input_array, output_array2)", "from __main__ import row_sum_gu, input_array, output_array2", number=100))

这给了我以下输出（时间确实有所不同）：

jit time: 12.04114792868495
guvectorize time: 5.415564753115177

因此，并行代码也仅仅快两倍（仅当行数是 CPU 核心数的整数倍时，否则性能优势就会减弱），即使它利用了所有 cpu 核心，而 jit 代码仅使用一个（使用 htop 进行验证）。

我在一台配备 4 个 AMD Opteron 6380 CPU（总共 64 个内核）、256 GB RAM 和 Red Hat 4.4.7-1 操作系统的机器上运行该程序。 我将 Anaconda 4.2.0 与 Python 3.5.2 和 Numba 0.26.0 结合使用。

我怎样才能进一步提高并行性能或者我做错了什么？

谢谢您的回答。

那是因为np.sum太简单了。使用 sum 处理数组不仅受到 CPU 的限制，还受到“内存访问”时间。所以投入更多的核心并不能解决问题much的差异（当然这取决于与 CPU 相关的内存访问速度）。

仅用于可视化np.sum是这样的（忽略除data):

def sum(data):
    sum_ = 0.
    data = data.ravel()
    for i in data.size:
        item = data[i]   # memory access (I/O bound)
        sum_ += item     # addition      (CPU bound)
    return sum

因此，如果大部分时间都花在访问内存上，那么如果将其并行化，您将不会看到任何真正的加速。但是，如果 CPU 密集型任务是瓶颈，那么使用更多内核将显着加快代码速度。

例如，如果您包含一些比加法慢的操作，您将看到更大的改进：

from math import sqrt
from numba import njit, jit, guvectorize
import timeit
import numpy as np

@njit
def square_sum(arr):
    a = 0.
    for i in range(arr.size):
        a = sqrt(a**2 + arr[i]**2)  # sqrt and square are cpu-intensive!
    return a

@guvectorize(["void(float64[:], float64[:])"], "(n) -> ()", target="parallel", nopython=True)
def row_sum_gu(input, output) :
    output[0] = square_sum(input)

@jit(nopython=True)
def row_sum_jit(input_array, output_array) :
    m, n = input_array.shape
    for i in range(m) :
        output_array[i] = square_sum(input_array[i,:])
    return output_array

I used IPythontimeit https://ipython.org/ipython-doc/3/interactive/magics.html#magic-timeit在这里，但它应该是等效的：

rows = int(64)
columns = int(1e6)

input_array = np.random.random((rows, columns))
output_array = np.zeros((rows))

# Warmup an check that they are equal 
np.testing.assert_equal(row_sum_jit(input_array, output_array), row_sum_gu(input_array, output_array2))
%timeit row_sum_jit(input_array, output_array.copy())  # 10 loops, best of 3: 130 ms per loop
%timeit row_sum_gu(input_array, output_array.copy())   # 10 loops, best of 3: 35.7 ms per loop

我只使用 4 个核心，因此这非常接近可能的加速极限！

请记住，并行计算只能显着加快计算速度如果作业受 CPU 限制.