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 不是为此而设计的。您看到的几乎所有执行时间都来自 wrapping/unwrapping 参数的开销、类型检查、Python 函数包装、引用计数等。此外使用的好处Numba 在这里非常小，因为 Numba 几乎没有优化 unicode 字符串操作。

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

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

普通函数和原始版本在我的机器上都需要 1.34 微秒。

此外，您可以反汇编 Numba 函数以查看执行一次调用需要执行多少指令，并深入了解开销来自何处。

如果您希望 Numba 有用，您需要 在编译函数中添加更多工作，可能是 直接在 arrays/lists[上工作]。如果由于输入类型的动态特性而无法做到这一点，那么 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 时）。

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

Performance loss in numba compiled logic comparison

python

compiler-construction

performance

numba