如何使两个数组连续以便 Numba 可以加速 np.dot()

我有以下代码：

import numpy as np
from numba import jit

Nx = 15
Ny = 1000

v = np.ones((Nx,Ny))
v = np.reshape(v,(Nx*Ny))
A = np.random.rand(Nx*Ny,Nx*Ny,5)
B = np.random.rand(Nx*Ny,Nx*Ny,5)
C = np.random.rand(Nx*Ny,5)
   
@jit(nopython=True)
def dotplus(B, v, C):
    return np.dot(B, v) + C

k = 2
D = dotplus(B[:,:,k], v, C[:,k])

我收到以下警告，我猜它指的是数组B[:,:,k] and v:

NumbaPerformanceWarning: np.dot() is faster on contiguous arrays, called on (array(float64, 2d, A), array(float64, 1d, C))
  return np.dot(B, v0) + C

有没有办法让两个数组连续，这样Numba就可以加速代码？

PS，如果您想知道其含义k，请注意这只是 MRE。在实际代码中，dotplus在一个内部被多次调用for循环不同的值k（因此，不同的切片B and C). The for循环更新的值v, but B and C不要改变。

缺陷是正确的。B[..., k]返回一个np.view() into B，但实际上并不复制任何数据。在内存中，视图的两个相邻元素的距离为B.strides[1]，其评估结果为B.shape[-1]*B.itemsize并且大于B.itemsize。因此，您的数组不是连续的。

最好的优化是向量化dotplus循环并写入

D = np.tensordot(B, v, axes=(1, 0)) + C

第二个最佳优化是重构，让批量维度成为数组的第一个维度。这可以在上述矢量化的基础上完成，并且通常是可取的。它看起来像

A = np.random.rand(5, Nx*Ny,Nx*Ny)
# rather than
A = np.random.rand(Nx*Ny,Nx*Ny,5)

如果您无法重构代码，则需要开始分析。您可以通过以下方式轻松临时交换轴

B = np.moveaxis(B, -1, 0)
some_op(B[k, ...], ...)
B = np.moveaxis(B, 0, -1) 

与 max9111 的评论相反，与np.ascontiguousarray()因为在这两种情况下都必须复制数据。也就是说，副本是O(Nx*Ny*k)+ 缓冲区分配。直接矩阵向量乘法是O(Nx*Ny)但你必须先收集元素，这确实很昂贵。这取决于您的特定架构和具体问题，因此分析是最佳选择。