为什么 numba 在使用字典时慢 50%

Question

我对 numba 使用字典时的性能感兴趣。我做了以下实验：

from numpy.random import randint
import numba as nb

@nb.njit
def foo_numba(a, b, c):
    N = 10**2
    d = {}
    for i in range(N):
        d[(randint(N), randint(N), randint(N))] = (a, b, c)
    return d


@nb.njit
def test_numba(numba_dict):
    s = 0
    for k in numba_dict:
        s += numba_dict[k][2]
    return s



def foo(a, b, c):
    N = 10**2
    d = {}
    for i in range(N):
        d[(randint(N), randint(N), randint(N))] = (a, b, c)
    return d



def test(numba_dict):
    s = 0
    for k in numba_dict:
        s += numba_dict[k][2]
    return s

a = randint(10, size=10)
b = randint(10, size=10)
c = 1.3

t_numba = foo_numba(a, b, c)
%timeit test_numba(t_numba)
t = foo(a, b, c)
%timeit test(t) 

令我惊讶的是，我得到的输出是：

9.28 µs ± 142 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
5.8 µs ± 375 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

为什么 numba 代码比 cPython 代码慢这么多？我是否错误地使用了 numba？

如果您执行 a = tuple(a) 或 b = tuple(b) ，那么速度减慢就会消失（似乎没有必要同时执行这两个操作！）。

Answer 1

正如 Jerome 在评论中提到的，这里最大的问题是您将 JIT 编译的时间作为基准测试的一部分。

只需在基准测试开始之前添加

test_numba(foo_numba(a, b, c))

将确保在测试之前编译内容，并消除大约一半的差异。

除此之外，您的测试用例并不是特别有用。 Python 的

numpy

s 已经进行了深度优化，C 层代码无法做任何有意义的事情来进一步优化访问，这与

dict

s/

list

s 不同（其中 C 层代码可以直接获取原始存储）内存访问避免了在 Python 级别执行相同工作时涉及的多层内部函数调用）或

tuple

s/

int

s（通过直接提取原始 C 值并执行 C 级别数学可以再次更有效地处理开销更少。理论上，您使用

float

数组

可以从

numpy中受益，但是您使用它们的方式

意味着这些好处并不存在。具体来说，在测试函数中

numba类型的only用途是： 通过迭代

numpy

加载
1. tuple

    中的 

   int

   （如果您只要求一个值，而不是 

   dict

    标量，

   numpy

    的 

   randint

    将返回 Python 

   int

   ）； 

   numpy

   迭代已经深度优化了，

   dict

   没办法

   再次抬头在

   numba
中说
2. tuple

   （再次，

   dict

   无能为力）

   从生成的 

   numba
数组中提取
3. single 值（numpy

    可以直接访问它，Python 需要构造一个包装对象，但这在 Python 中进行了大量优化，并且 

   numba

    的直接访问收益非常大）当无法进行矢量化时，第一次访问的成本差异不够大，并且您不会再次访问同一数组，因此不会产生大量工作收益）

   将其添加到

   numba
。这是
4. scouldnumbashelp with, if it recognizes

   int

   is an

   +=

   that never overflows a 64 bit value. But again, this is something the most modern versions of Python have optimized quite a bit (by 3.12, the bytecode self-modifies to handle cases where the

   int`s 的一件事，并且绕过了旧 Python 必须支付的大部分开销）。

   简而言之：

always involves

1. 永远不会

    不会为 

   numbas 做太多事情 当您进行大量简单数学运算时，dict
的增益效果最佳，最好是使用
2. numba

    数组进行 

   数学运算。从

   numpy

   数组的

   dict

   中提取单个值无法克服

   numpy

   在

   尝试

   优化困难事物时所付出的成本，这些成本被浪费在简单的事情上。