并行化numba循环

Question

之前，我问过a question一个相对简单的循环，Numba无法并行化。结果证明，一种解决方案可以使所有循环都显式化

现在，我需要做相同任务的一个简单版本：我现在有数组alpha和beta，分别是(m,n)和(b,m,n),，我想计算参数的2D切片的Frobenius乘积，并找到使该乘积最大化的β切片。以前，有一个额外的、大的第一维度alpha，所以我是在这个维度上并行化的；现在我想在beta的第一维上进行并行化，因为当b>；1000.

如果我天真地修改了用于前一个问题的代码，我将获得：

@njit(parallel=True)
def parallel_value_numba(alpha,beta):
    dot = np.zeros(beta.shape[0])
        for i in prange(beta.shape[0]):
            for j in prange(beta.shape[1]):
                for k in prange(beta.shape[2]):
                    dot[i] += alpha[j,k]*beta[i, j, k]
        index=np.argmax(dot)
        value=dot[index]
    return value,index

但出于某种原因，Numba不喜欢这样，并抱怨：

numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
scalar type memoryview(float64, 2d, C) given for non scalar argument #3

所以我试着

@njit(parallel=True)
def parallel_value_numba_2(alpha,beta):
    product=np.multiply(alpha,beta)
    dot1=np.sum(product,axis=2)
    dot2=np.sum(dot1,axis=1)
    index=np.argmax(dot2)
    value=dot2[index]
    return value,index

只要在传递给函数之前将alpha广播到beta.shape，并且原则上Numba能够并行化numpy操作，它就会编译。但是它运行得非常慢，比串行的纯Python代码慢得多

def einsum_value(alpha,beta):
    dot=np.einsum('kl,jkl->j',alpha,beta)
    index=np.argmax(dot)
    value=dot[index]
    return value,index

所以，我当前的工作代码使用了最后一个实现，但是这个函数仍然是运行时的瓶颈，我想加快它。有人能说服Numba以可观的速度并行化这个函数吗

Answer 1

这并不完全是一个解决方案，但格式化注释更难

Numba根据传递给函数的参数生成不同的代码。例如，您的代码使用以下示例：

>>> alpha = np.random.random((5, 4))
>>> beta = np.random.random((3, 5, 4))
>>> parallel_value_numba(alpha, beta)
(5.89447648574048, 0)

为了诊断问题，有必要提供导致问题的特定参数值的示例

读取错误消息时，您似乎正在传递一个memoryview对象，但Numba可能没有完整的support

作为旁注，您不需要在每个循环中使用prange。在外循环中使用它通常就足够了，只要预期迭代的数量大于机器中的内核数量