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

Question

我有以下代码：

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]和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多次调用。循环更新 的值，但和不改变。forkBCforvBC

Answer 1

缺陷是正确的。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)，但是你必须先收集元素，这非常昂贵。这取决于您的特定架构和具体问题，因此分析是最佳选择。