Numba 慢速数组元素分配给变量

Question

这是一个人为的测试用例，但希望它足以传达要点并提出问题。在 Numba njit 函数内部，我注意到将本地计算的值分配给数组元素非常昂贵。以下是两个示例函数：

from numba import njit
import numpy as np

@njit
def slow_func(x, y):
    result = y.sum()
    
    for i in range(x.shape[0]):
        if x[i] > result:
            x[i] = result
        else:
            x[i] = result

@njit
def fast_func(x, y):
    result = y.sum()
    
    for i in range(x.shape[0]):
        if x[i] > result:
            z = result
        else:
            z = result

if __name__ == "__main__":
    x = np.random.rand(100_000_000)
    y = np.random.rand(100_000_000)

    %timeit slow_func(x, y)  # 177 ms ± 1.49 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
    %timeit fast_func(x, y)  # 407 ns ± 12.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

我知道这两个函数做的事情并不完全一样，但我们暂时不用担心，继续专注于“慢任务”。此外，由于 Numba 的延迟初始化，上述时间已在 JIT 编译后重新运行。请注意，这两个函数都将result 分配给x[i] 或z，并且在这两种情况下分配的数量相同。但是，将result 分配给z 的速度要快得多。有没有办法让slow_func 和fast_func 一样快？

Answer 1

正如@PaulPanzer 已经指出的那样，你的快速函数一旦优化就什么都不做 - 所以你看到的基本上是调用 numba 函数的开销。

有趣的是，为了进行这种优化，numba 必须将 np.sum 替换为自己的 sum-implementation - 否则优化器将无法放弃对该函数的调用，因为它无法查看np.sum 的实现，必须假设调用此函数会产生副作用。

让我们只测量 numba 的总和：

from numba import njit
@njit
def only_sum(x, y):
    return y.sum()

%timeit only_sum(y,x) 
# 112 ms ± 623 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

。 好吧，这令人失望：我知道我的机器每秒可以进行超过 10^9 的加法运算，并且可以从 RAM 读取高达 13GB/s 的速度（大约有 0.8GB 数据，因此它不适合缓存），这意味着我希望总和在 60-80 毫秒之间使用。

如果我使用 numpy 的版本，它确实可以：

%timeit y.sum()
# 57 ms ± 444 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

这听起来很对！我假设，numba 不使用成对加法，因此是 slower（如果 RAM 足够快成为瓶颈）并且比 numpy 的版本少 precise。

如果我们只看值的写法：

@njit
def only_assign(x, y):
    res=y[0]
    for i in range(x.shape[0]):
        x[i]=res

%timeit only_assign(x,y)
85.2 ms ± 417 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

所以我们看到它真的比阅读慢。这个great answer解释了其原因（以及如何修复它）：numba（对吗？）没有绕过的缓存更新。

---

简而言之：虽然在 numba 中分配值并不是很慢（即使它可以通过使用 non-temporal memory accesses 来加速），但真正慢的部分是求和（似乎没有使用成对求和) - 它不如 numpy 的版本。