Python Numba/jit 条件和递归(堆栈)使用

Question

全部，

我正在使用 numba JIT 来加速我的 Python 代码，但即使没有安装 numba 和 LLVM，代码也应该可以正常运行。

我的第一个想法是按如下方式进行:

use_numba = True
try:
    from numba import jit, int32
except ImportError, e:
    use_numba = False

def run_it(parameters):
    # do something
    pass

# define wrapper call function with optimizer
@jit
def run_it_with_numba(parameters):
    return run_it(parameters)

# [...]
# main program 
t_start = timeit.default_timer()

# this is the code I don't like 
if use_numba:
    res = run_it_with_numba(parameters)
else:
    res = run_it(parameters)

t_stop = timeit.default_timer()
print "Numba: ", use_numba, " Time: ", t_stop - t_start

这并没有像我预期的那样工作，因为编译似乎只适用于 run_it_with_numba() 函数——它基本上什么都不做——但不适用于从该函数调用的子例程。

只有当我在包含工作负载的函数上应用 @jit 时，结果才会变得更好。

有没有机会在主程序中避免包装函数和 if 子句？

有没有办法告诉 Numba 优化从我的入口函数调用的子例程？因为 run_it() 还包含一些函数调用，我希望 @jit 能够处理它。

铜，麦酒

Answer 1

如果没有安装 Numba，你可以提供一个什么都不做的 jit 版本:

use_numba = True
try:
    from numba import jit, int32
except ImportError, e:
    use_numba = False
    from _shim import jit, int32

@jit
def run_it(parameters):
    # do something
    pass

# [...]
# main program 
t_start = timeit.default_timer()

res = run_it(eval(row[0]), workfeed, instrument)

t_stop = timeit.default_timer()
print "Numba: ", use_numba, " Time: ", t_stop - t_start

_shim.py 只包含:

def jit(*args, **kwargs):
    def wrapper(f):
        return f
    if len(args) > 0 and (args[0] is marker or not callable(args[0])) \
        or len(kwargs) > 0:
        # @jit(int32(int32, int32)), @jit(signature="void(int32)")
        return wrapper
    elif len(args) == 0:
        # @jit()
        return wrapper
    else:
        # @jit
        return args[0]

def marker(*args, **kwargs): return marker

int32 = marker