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 被调用时）。