含可选返回数组的Numba函数执行变慢问题及优化咨询

我尝试编写一个可选择性返回内部计算数组的Numba函数，发现即便未触发返回数组的逻辑，该函数的执行速度也比无此选项的同功能函数慢。以下是示例代码：

import numpy as np
from numba import jit

@jit(nopython=True, nogil=True, cache=True)
def testfun1(array_in, arrays_out=np.empty((0,0), np.float64)):
    array_add_1 = array_in + 1
    array_add_2 = array_in + 2
    array_sum_1 = sum(array_add_1)
    array_sum_2 = sum(array_add_2)
    
    if arrays_out.shape == (10,len(array_in)):
        # write to an array
        arrays_out[0,:] = array_add_1
        arrays_out[1,:] = array_add_1
        arrays_out[2,:] = array_add_1
        arrays_out[3,:] = array_add_1
        arrays_out[4,:] = array_add_1
        
        arrays_out[5,:] = array_add_2
        arrays_out[6,:] = array_add_2
        arrays_out[7,:] = array_add_2
        arrays_out[8,:] = array_add_2
        arrays_out[9,:] = array_add_2
    
        
    return (array_sum_1, array_sum_2)

@jit(nopython=True, nogil=True, cache=True)
def testfun2(array_in, arrays_out=np.empty((0,0), np.float64)):
    array_add_1 = array_in + 1
    array_add_2 = array_in + 2
    array_sum_1 = sum(array_add_1)
    array_sum_2 = sum(array_add_2)
    
    if arrays_out.shape == (10,len(array_in)):
        array_add_1 = 0 # do something
        array_add_2 = 0 # do something
        
    return (array_sum_1, array_sum_2)


array_in = np.arange(0,10000)

计时结果：

%timeit testfun1(array_in) # 70.1 µs ± 346 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit testfun2(array_in) # 69.6 µs ± 217 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

注：这是简化示例，实际代码中因数组更多更大，时间差异更显著。

明明未传入符合尺寸的arrays_out，两个函数逻辑一致，为何testfun1更慢？有没有能实现可选返回数组且不影响无返回场景执行效率的替代方案？

为什么testfun1更慢？ #

Numba的nopython模式在编译阶段生成机器码，虽然你的if分支未执行，但testfun1中的大量数组赋值操作会被纳入编译优化逻辑：

• Numba可能为这些数组操作提前分配临时寄存器或预留内存相关指令，即便分支不触发，这些编译产生的额外逻辑会带来微小运行开销（大数组场景下差异被放大）。
• testfun1中array_add_1和array_add_2需保持数组形态以支持赋值，而testfun2分支内将其赋值为0，Numba可对后者进行更激进的优化（比如提前释放数组引用、简化内存管理）。

替代方案 #

1. 拆分函数（推荐） #

将核心计算与数组输出逻辑拆分为独立函数，用户按需调用：

import numpy as np
from numba import jit

@jit(nopython=True, nogil=True, cache=True)
def core_calc(array_in):
    array_add_1 = array_in + 1
    array_add_2 = array_in + 2
    array_sum_1 = sum(array_add_1)
    array_sum_2 = sum(array_add_2)
    return array_sum_1, array_sum_2, array_add_1, array_add_2

@jit(nopython=True, nogil=True, cache=True)
def with_output(array_in, arrays_out):
    sum1, sum2, add1, add2 = core_calc(array_in)
    if arrays_out.shape == (10, len(array_in)):
        for i in range(5):
            arrays_out[i,:] = add1
        for i in range(5,10):
            arrays_out[i,:] = add2
    return sum1, sum2

# 无输出需求调用
%timeit core_calc(array_in)[:2]  # 仅取求和结果
# 有输出需求调用
out_arr = np.empty((10, len(array_in)), np.float64)
%timeit with_output(array_in, out_arr)

这种方式完全隔离两种场景的编译逻辑，无输出场景不受数组赋值代码影响，性能最优。

2. 使用Numba重载功能 #

利用@overload根据输入参数生成不同编译版本，保持接口统一：

from numba import overload, jit, types

def optional_output_func(array_in, arrays_out=None):
    pass

@overload(optional_output_func)
def _overload_optional_output(array_in, arrays_out=None):
    if arrays_out is None:
        # 无输出参数版本
        def impl(array_in, arrays_out=None):
            array_add_1 = array_in + 1
            array_add_2 = array_in + 2
            return sum(array_add_1), sum(array_add_2)
        return impl
    else:
        # 有输出参数版本
        def impl(array_in, arrays_out):
            array_add_1 = array_in + 1
            array_add_2 = array_in + 2
            sum1 = sum(array_add_1)
            sum2 = sum(array_add_2)
            if arrays_out.shape == (10, len(array_in)):
                for i in range(5):
                    arrays_out[i,:] = array_add_1
                for i in range(5,10):
                    arrays_out[i,:] = array_add_2
            return sum1, sum2
        return impl

# 包装为jit函数
@jit(nopython=True, nogil=True, cache=True)
def wrapper(array_in, arrays_out=None):
    return optional_output_func(array_in, arrays_out)

# 无输出调用
%timeit wrapper(array_in)
# 有输出调用
out_arr = np.empty((10, len(array_in)), np.float64)
%timeit wrapper(array_in, out_arr)

Numba会根据输入参数自动匹配对应编译版本，无输出场景性能不受影响。

3. 提前分支判断优化 #

在函数开头明确判断是否需要输出，让Numba对不同分支独立优化：

@jit(nopython=True, nogil=True, cache=True)
def optimized_testfun(array_in, arrays_out=np.empty((0,0), np.float64)):
    need_output = (arrays_out.shape == (10, len(array_in)))
    
    array_add_1 = array_in + 1
    array_add_2 = array_in + 2
    array_sum_1 = sum(array_add_1)
    array_sum_2 = sum(array_add_2)
    
    if need_output:
        for i in range(5):
            arrays_out[i,:] = array_add_1
        for i in range(5,10):
            arrays_out[i,:] = array_add_2
    
    return (array_sum_1, array_sum_2)

这种方式让Numba能识别need_output=False的分支，优化掉所有数组赋值相关逻辑，减少不必要的编译开销。

内容的提问来源于stack exchange，提问作者Scooba