python - 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。编写 native C/C++ 模块可能会有所帮助，但大部分时间将花在操作动态对象和 unicode 字符串以及进行类型自省(introspection)上，除非您重写整个代码。

更新
上述开销仅在从 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 时)。