numba编译逻辑比较中的性能损失

Question

以下用于逻辑比较的 numba 编译函数中性能下降的原因可能是什么：

from numba import njit

t = (True, 'and_', False)

#@njit(boolean(boolean, unicode_type, boolean))    
@njit
def f(a,b,c):
    if b == 'and_':
        out = a&c
    elif b == 'or_':
        out = a|c
    return out
x = f(*t)
%timeit f(*t)
#1.78 µs ± 9.52 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

---

%timeit f.py_func(*t)
#108 ns ± 0.0042 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

要按照答案中的建议进行大规模测试：

x = np.random.choice([True,False], 1000000)
y = np.random.choice(["and_","or_"], 1000000)
z = np.random.choice([False, True], 1000000)

#using jit compiled f
def f2(x,y,z):
    L = x.shape[0]
    out = np.empty(L)
    for i in range(L):
        out[i] = f(x[i],y[i],z[i])
    return out

%timeit f2(x,y,z)
#2.79 s ± 86.4 ms per loop

---

#using pure Python f
def f3(x,y,z):
    L = x.shape[0]
    out = np.empty(L)
    for i in range(L):
        out[i] = f.py_func(x[i],y[i],z[i])
    return out

%timeit f3(x,y,z)
#572 ms ± 24.3 ms per

我是否遗漏了一些东西，是否有办法编译“快速”版本，因为这将成为执行约 1e6 次循环的一部分。

Answer 1

您的工作粒度太小。 Numba 不是为此而设计的。您看到的几乎所有执行时间都来自包装/展开参数、类型检查、Python 函数包装、引用计数等的开销。此外，使用 Numba 的好处非常小，因为 Numba 几乎不优化 unicode 字符串操作。

检验这一假设的一种方法是执行以下简单函数：

@njit
def f(a,b,c):
    return a
x = f(True, 'and_', False)
%timeit f(True, 'and_', False)

在我的机器上，这个简单函数和原始版本都需要 1.34 µs。

此外，您还可以

反汇编 Numba 函数，查看仅执行一次调用就执行了多少指令，并深入了解开销来自何处。

如果您希望 Numba 有用，您需要在编译函数中

添加更多工作，可能是通过直接处理数组/列表。如果由于输入类型的动态特性而无法实现这一点，那么 Numpy 可能不是适合此操作的工具。您可以尝试修改一下代码并使用 PyPy 代替。 编写本机 C/C++ 模块可能会有所帮助，但大多数时间将花费在操作动态对象和 unicode 字符串以及进行类型自省上，除非您重写整个代码。

---

更新

仅当从 Python 类型转换到 Numba 时（反之亦然），才会支付上述开销。您可以通过以下基准看到这一点：

@njit
def f(a,b,c):
    if b == 'and_':
        out = a&c
    elif b == 'or_':
        out = a|c
    return out
@jit
def manyCalls(a, b, c):
    res = True
    for i in range(1_000_000):
        res ^= f(a, b, c ^ res)
    return res

t = (True, 'and_', False)
x = manyCalls(*t)
%timeit manyCalls(*t)

在我的机器上，

调用

manyCalls

 需要 3.62 毫秒。这意味着每次调用 

f

 平均需要 3.6 ns（16 个周期）。这意味着开销只需支付一次（当调用 

manyCalls

 时）。