"如何在 Numba 中强制循环提升？"

Question

我正在尝试实现numpy.take()<\/code>的 Numba 编译版本，但我对 Numba 感到非常困惑。

首先，据我了解，在ndarray<\/code>模式下无法在 numba 函数内创建新的nopython<\/code> 。 当前版本的文档似乎没有在任何地方提及它，但我能够在旧的 Numba v0.15 文档<\/a>中找到它。 所以我不得不放弃nopython<\/code>模式。

其次，我不明白循环提升是如何工作的，至少在我的测试中是这样。 这是我的代码：

from numba import jit
import numpy as np
import time

@jit(forceobj = True)
def _take1(arr, idxs):
    res = np.ndarray((idxs.size,), arr.dtype)
    lastIdx = arr.size - 1
    for i in range(idxs.size):
        idx = idxs[i]
        if idx > lastIdx:
            idx = lastIdx
        elif idx < 0:
            idx = 0
        res[i] = arr[idx]
    return res

def _take2(arr, idxs):
    res = np.ndarray((idxs.size,), arr.dtype)
    lastIdx = arr.size - 1
    for i in range(idxs.size):
        idx = idxs[i]
        if idx > lastIdx:
            idx = lastIdx
        elif idx < 0:
            idx = 0
        res[i] = arr[idx]
    return res

sz = 2000000
arr = np.arange(sz, dtype = np.int32)
idxs = np.arange(sz, 0, -1, dtype = np.int32)

start = time.time()
_take2(arr, idxs)
end = time.time()
print("Elapsed (plain python)      = %s" % (end - start))

start = time.time()
_take1(arr, idxs)
end = time.time()
print("Elapsed (with compilation)  = %s" % (end - start))

start = time.time()
_take1(arr, idxs)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

Answer 1

我不能在 nopython 模式下用 Numba 创建一个新的 ndarray 是真的吗？

<\/blockquote>

这是错误<\/strong>的。 以下工作代码执行此操作：

       

        为什么回退到对象模式的@jit 比@jit(forceobj = True) 工作得快很多？
<\/blockquote>
 
你应该避免像瘟疫一样的对象模式<\/strong>。 它通常非常低效并且通常不是真正有用，除非您处理由于 GIL、引用计数、分配、动态类型和许多使 CPython 变慢的性能问题而导致计算缓慢的 Python 对象。 该文档<\/a>指出：
        
 
        
 
         
虽然在对象模式下使用循环提升可以提高一些性能，但让函数在非 Python 模式下<\/strong>编译确实是获得良好性能的关键。
         
         
 
         
因此，挑战在于使用njit<\/code>模式构建您的代码。 在您的情况下，这很简单：您只需要使用np.empty<\/code>和njit<\/code> 。
         
 
         
这使我的机器上的代码非常熟练：
请注意，您可以提供_take1<\/code>的签名，以便 Numba 可以急切地<\/a>编译它。

                     
              
 
               
 Numba 中是否有没有循环提升的对象模式？ 如果是这样，它的目的是什么？
               
               
 
               
对象模式的目的只是为了与纯 Python 代码兼容。 当 Numba 代码与纯 Python 函数混合时，这很有用。 纯 Python 函数调用会很慢，但至少 Numba 可以在 Numba 代码中间调用它们。 对象模式在优化的 Numba 代码中通常不是很有用，如果使用它，通常需要 looplifting 以提高热循环的性能。 否则，在目标情况下使用 Numba 肯定不是很有用（罕见的例外包括包装代码）。