Numba JIT 急切编译无法按预期工作

如何解决Numba JIT 急切编译无法按预期工作

我正在尝试学习 Numba 来加速我的 Python 代码，我打算使用 eager compilation mode of Numba JIT。

常用的JIT模式可以通过如下代码实现：

@jit(nopython=True)
def jit_go_fast_lazy(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i,i])
    return a + trace

当运行上面的程序时，jit_go_fast()的第一次执行会很慢，因为编译需要很多时间，但是之后的执行速度会非常快。

由于我的程序对执行时间很敏感，所以必须减少开销编译时间，所以我打算选择eager编译模式，根据official tutorial应该是这样的：

@jit(float64[:,:](int64[:,:]),nopython=True)
def jit_go_fast_eager(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i,i])
    return a + trace

（这里的代码示例修改自Numba official site）

在测试上面代码的运行时间时，我使用如下代码：

x = np.arange(100).reshape(10,10)

start = time.time()
jit_go_fast_lazy(x)
end = time.time()
print("Elapsed = %s" % (end - start))

start = time.time()
jit_go_fast_lazy(x)
end = time.time()
print("Elapsed = %s" % (end - start))

start = time.time()
jit_go_fast_eager(x)
end = time.time()
print("Elapsed = %s" % (end - start))

start = time.time()
jit_go_fast_eager(x)
end = time.time()
print("Elapsed = %s" % (end - start))

在使用 Python 3.9.1 的 macOS 上运行此测试时，结果如下：

First run of JIT lazy compilation: 0.14231 sec
Second run of JIT lazy compilation: 3.0994e-06 sec

First run of JIT eager compilation: 8.8930e-05 sec
Second run of JIT eager compilation: 1.9073e-06 sec

据我所知，jit_go_faster_eager() 的运行时间应该与我第一次在我的程序中运行它的时间大致相同（因为它已经编译过）。但是，结果显示，当第二次执行 jit_go_faster_eager() 时，速度大约快了 10 倍。

我对这个令人困惑的结果进行了一些搜索，并找到了 this question。我将 jit_go_fast_eager() 的装饰器更改为：

@jit(float64[:,::1](int64[:,::1]),nopython=True)

结果是：

First run of JIT lazy compilation: 0.14407 sec
Second run of JIT lazy compilation: 3.0994e-06 sec

First run of JIT eager compilation: 6.0081e-05 sec
Second run of JIT eager compilation: 2.1458e-06 sec

运行时间只有很小的变化。

有人可以对此提供一些帮助吗？我的主要问题是：

• 为什么在第一次和第二次执行 JIT 急切编译函数时会出现巨大的运行时间差异？
• 加速此类功能的最佳做法是什么？虽然第二次执行总是要快得多，但第一次执行可能会干扰程序的整个过程。

谢谢！