Numba 没有提高性能

Question

我正在测试一些采用 numpy 数组的函数的 numba 性能，并进行比较：

import numpy as np
from numba import jit, vectorize, float64
import time
from numba.core.errors import NumbaWarning
import warnings

warnings.simplefilter('ignore', category=NumbaWarning)

@jit(nopython=True, boundscheck=False) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a):     # Function is compiled to machine code when called the first time
    trace = 0.0
    for i in range(a.shape[0]):   # Numba likes loops
        trace += np.tanh(a[i, i]) # Numba likes NumPy functions
    return a + trace              # Numba likes NumPy broadcasting
   
class Main(object):
    def __init__(self) -> None:
        super().__init__()
        self.mat     = np.arange(100000000, dtype=np.float64).reshape(10000, 10000)

    def my_run(self):
        st = time.time()
        trace = 0.0
        for i in range(self.mat.shape[0]):   
            trace += np.tanh(self.mat[i, i]) 
        res = self.mat + trace
        print('Python Diration: ', time.time() - st)
        return res                           
    
    def jit_run(self):
        st = time.time()
        res = go_fast(self.mat)
        print('Jit Diration: ', time.time() - st)
        return res
        
obj = Main()
x1 = obj.my_run()
x2 = obj.jit_run()

输出是：

Python Diration:  0.2164750099182129
Jit Diration:  0.5367801189422607

如何获得此示例的增强版本？

Answer 1

Numba 实现较慢的执行时间是由于编译时间，因为 Numba 在函数使用时编译函数（除非参数类型发生变化，仅在第一次编译）。它这样做是因为它在调用函数之前无法知道参数的类型。希望您可以为 Numba 指定参数类型，以便它可以直接编译函数（在执行装饰器函数时）。这是生成的代码：

@njit('float64[:,:](float64[:,:])')
def go_fast(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

请注意，njit 是jit+nopython=True 的快捷方式，并且boundscheck 默认已设置为False（请参阅doc）。

在我的机器上，这导致 Numpy 和 Numba 的执行时间相同。实际上，执行时间不受tanh 函数计算的限制。它以表达式a + trace 为界（适用于 Numba 和 Numpy）。预计执行时间相同，因为两者都以相同的方式实现：它们创建一个临时的新数组来执行加法。由于page faults 和RAM 的使用（a 完全从 RAM 中读取并且临时数组完全存储在 RAM 中），创建一个新的临时数组的成本很高。如果您想要更快的计算，那么您需要执行 in-place 操作（这可以防止在 x86 平台上出现页面错误和昂贵的cache-line write allocations）。