python - 如何在多个功能中最佳地使用 numba？

Question

假设我有两个功能

def my_sub1(a):
    return a + 2

def my_main(a):
    a += 1
    b = mysub1(a)
    return b

我想使用像 Numba 这样的即时编译器使它们更快。这是否会比我将所有内容重构为一个函数要慢

def my_main(a):
    a += 1
    b = a + 2
    return b

因为 Numba 可以在第二种情况下进行更深入的优化？当然，我的实际功能要复杂得多。
如果 my_sub1，整个情况也会变得更加困难。函数 get 被多次调用 - 重构(和维护将成为拖累)。 Numba 如何解决这个问题？

Answer 1

Tl;博士: Numba 能够内联其他 Numba 函数，并且仅在使用 native 类型时才执行相对高级的过程间优化(在这种情况下两个函数同样快)，但不适用于 Numpy 数组。

您可以分析 Numba 生成的结果汇编代码，以检查这两个函数是如何优化的。这是一个带有整数的示例:

import numba as nb

@nb.njit('int64(int64)')
def my_sub1(a):
    return a + 2

@nb.njit('int64(int64)')
def my_main(a):
    a += 1
    b = my_sub1(a)
    return b

open('my_sub1.asm', 'w').write(list(my_sub1.inspect_asm().values())[0])
open('my_main.asm', 'w').write(list(my_main.inspect_asm().values())[0])

这将生成两个程序集文件。如果你比较这两个文件，你会发现唯一的实际区别(除了不同的名称)是第一个做 addq $2, %rdx而第二个做 addq $3, %rdx .这意味着 Numba 成功地将调用内联到 my_sub1在 my_main并合并总和。这是汇编代码的重要部分:

_ZN8__main__12my_sub1$2413Ex:
    addq    $2, %rdx
    movq    %rdx, (%rdi)
    xorl    %eax, %eax
    retq

_ZN8__main__12my_main$2414Ex:
    addq    $3, %rdx
    movq    %rdx, (%rdi)
    xorl    %eax, %eax
    retq

对于 64 位浮点数，只要使用 fastmath=True，结果就相同。因为浮点加法不是结合的。
对于 Numpy 数组，生成的代码非常庞大，很难比较这两个代码。然而， my_sub1函数似乎不再内联，Numba 似乎无法合并 Numpy 计算(生成的代码中存在用于两个数组求和的两个不同的矢量化循环)。请注意，这与许多 C/C++ 编译器所做的类似。因此，最好在代码的性能关键部分自己内联函数。