动态编译 numba jit答案

【问题标题】：Dynamically compile numba jit动态编译 numba jit
【发布时间】：2021-03-17 15:06:58
【问题描述】：

我有许多函数，我寻求支持单线程和多线程执行，这使我的代码行数增加，产生冗余重复：

@jit(nopython=True, cache=True)
def fn1(x):
    for i in range(len(x)):
        # [code body]

@jit(nopython=True, cache=True, parallel=True)
def fn1_par(x):
    for i in prange(len(x)):
        # [code body]

两者的代码体完全相同，唯一不同的是@jit(parallel=True) 和range -> prange。有没有办法让我在 single 函数中表达这个逻辑，同时保留 cache=True 的好处？

注意：使用 'wrapper wrapper'（osvil 的回答）的解决方法因 cache=True 而失败；第一次调用的配置保持不变（例如，第一个函数的 parallel=False 将覆盖第二个函数的 parallel=True）。更改包装函数的 __name__ 无济于事。已打开Issue。

【问题讨论】：

无需将prange 替换为range。当使用parallel=False 调用时，numba 将自动将prange 替换为range，允许您使用@osvil 的答案中概述的方法

标签： python jit numba

【解决方案1】：

numba.jit 只是一个装饰器。您可以定义一次函数，然后在同一函数上多次“手动”应用装饰器。

在我的脑海中，它看起来像：

def python_function(x):
   ...

numba_nopython_function = numba.jit(nopython=True)(python_function)
numba_nopython_parallel_function = numba.jit(nopython=True, parallel=True)(python_function)

通过numba_nopython_function 使用该函数将使用普通的numba 版本。使用 numba_nopython_parallel_function 将使用 numba 并行函数。如果您愿意，还可以通过 python_function 本身访问原始 python 函数。它符合 DRY 原则。随意命名，因为我冗长的名字并不实用。

不过，我不知道这将如何与函数的缓存交互。

【讨论】：

@OverLordGoldDragon 如果你使用这种方式，你可以在定义函数时使用prange。这是可行的，因为当用parallel=False 调用时，numba 会自动将prange 替换为range。
对，够简单，会试试。 @Bob 尝试使用和不使用prange 的非确定性计算，会得到不同的结果；这表明range -> prange 不会自动发生。
@OverLordGoldDragon：你搞错了。 range 不会自动转换为 prange；那将是不安全的。 Bob 是说在非并行模式下，prange 是 range 的别名。
@user2357112supportsMonica 啊，我错了。
这不适用于cache=True；首先包装python_function 的配置将持续存在，所以如果例如parallel=False 是第一个然后parallel=True 是无效的。（所以缓存确实有效）。