将 numba 与 np.concatenate 并行使用效率不高？

Question

我在让 np.concatenate 高效并行化时遇到了一些麻烦。这是一个最小的工作示例。 （我知道在这里我可以分别对a 和b 求和，但我专注于并行化连接操作，因为这是我在项目中需要做的。然后我将对连接数组进行进一步的操作，例如作为排序。）

无论我在多少个内核上运行它，它似乎总是需要相同的时间（约 10 秒）。如果有的话，它在更多内核上会更慢。我尝试在装饰器中使用 nogil=True 在 cc 上释放 GIL，但无济于事。请注意，即使没有加速，在计算期间所有内核都在使用中。

谁能帮帮我？

非常感谢

from numba import prange, njit
import numpy as np


@njit()
def cc():

    r = np.random.rand(20)
    a = r[r < 0.5]
    b = r[r > 0.7]

    c = np.concatenate((a, b))

    return np.sum(c)


@njit(parallel=True)
def cc_wrap():
    n = 10 ** 7
    result = np.empty(n)
    for i in prange(n):
        result[i] = cc()

    return result

cc_wrap()

Answer 1

主要问题来自分配器的争用。

cc 函数创建许多隐式的小型临时数组。例如np.random.rand 和r < 0.5 甚至a = r[condition]，更不用说np.concatenate。临时数组通常需要使用给定的分配器在堆中分配。标准库提供的默认分配器不能保证使用多个线程很好地扩展。分配不能完美扩展，因为在分配数据的线程之间需要昂贵的隐式同步。例如，一个线程可以分配一个被另一个线程删除的数组。在最坏的情况下，分配器可以序列化分配/删除。由于分配数据的成本与对分配的数据执行的工作相比是巨大的，因此同步的开销占主导地位，并且整体执行是序列化的。实际上，更糟糕的是，因为连续时间已经被开销所支配。

请注意，积极优化的编译器可能会在 堆栈 上分配数组，因为它们不会 escape the function。然而，不幸的是，Numba 并没有明显地做到这一点。此外，假设 Numba 线程从不删除其他线程分配的数据，分配器可以调整为使用每个线程池很好地扩展（尽管我不完全确定，但可能是这种情况）。尽管如此，分配的内存池仍需要向通常不能很好扩展的操作系统（尤其是 Windows AFAIK）请求。

最好的解决方案就是避免创建隐式临时数组。这可以通过将per-worker local temporary array 与分区算法 结合使用来实现。请注意，可以通过将类型指定为 Numba，来提前编译函数。

这是最终的实现：

import numba as nb
import numpy as np
import random

@nb.njit('float64(float64[::1])')
def cc(tempBuffer):
    assert tempBuffer.size >= 20

    # View to the temporary scratch buffer
    r = tempBuffer[0:20]

    # Generate 20 random numbers without any allocation
    for i in range(20):
        r[i] = random.random()

    j = 0

    # Partition the array so to put values smaller than
    # a condition in the beginning.
    # After the loop, r[0:split] contains the filtered values.
    for i in range(r.size):
        if r[i] < 0.5:
            r[i], r[j] = r[j], r[i]
            j += 1

    split = j

    # Partition the rest of the array.
    # After the loop, r[split:j] contains the other filtered values.
    for i in range(j, r.size):
        if r[i] > 0.7:
            r[i], r[j] = r[j], r[i]
            j += 1

    # Note that extracting contiguous views it cheap as 
    # it does not create a new temporary array
    # a = r[0:split]
    # b = r[split:j]
    c = r[0:j]

    return np.sum(c)

@nb.njit('float64[:]()', parallel=True)
def cc_wrap():
    n = 10 ** 7
    result = np.empty(n)

    # Preallocate some space once for all threads
    globalTempBuffer = np.empty((nb.get_num_threads(), 64), dtype=np.float64)

    for i in nb.prange(n):
        threadId = nb.np.ufunc.parallel._get_thread_id()
        threadLocalBuffer = globalTempBuffer[threadId]
        result[i] = cc(threadLocalBuffer)

    return result

cc_wrap()

请注意，线程本地操作有点棘手，通常不需要。在这种情况下，仅使用分区算法来减少分配就可以看到加速。但是，由于临时数组的大小和分配的数量非常小，分配的开销仍然很大。

另请注意，此代码中并不严格要求 r 数组，因为随机数可以就地求和。这可能不符合您对实际用例的需求。这是一个（更简单的实现）：

@nb.njit('float64()')
def cc():
    s = 0.0
    for i in range(20):
        e = random.random()
        if e < 0.5 or e > 0.7:
            s += e
    return s

@nb.njit('float64[:]()', parallel=True)
def cc_wrap():
    n = 10 ** 7
    result = np.empty(n)
    for i in nb.prange(n):
        result[i] = cc()
    return result

cc_wrap()

这是我的 6 核机器上的计时：

# Initial (sequential):      8.1 s
# Initial (parallel):        9.0 s
# Array-based (sequential):  2.50 s
# Array-based (parallel):    0.41 s
# In-place (sequential):     1.09 s
# In-place (parallel):       0.19 s

最终，最快的并行版本比原始版本快 47 倍（并且几乎完美地扩展）。