如何在多个功能中最佳地使用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；dr:Numba能够内联其他Numba函数，并且它仅在使用本机类型时执行相对高级的过程间优化（在本例中，这两个函数的速度相同），但不使用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])

这将生成两个程序集文件。如果比较这两个文件，您将看到唯一的实际区别（除了不同的名称之外）是第一个doaddq $2, %rdx，而第二个doaddq $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++编译器所做的类似。因此，在代码的性能关键部分自己内联函数可能更好