如何使两个数组连续，以便 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,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],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 在 for 循环内针对不同的 k 值（因此，B 和 C 的不同切片）被多次调用。 for 循环更新 v 的值，但 B 和 C 不会改变。

解决方法

缺陷是正确的。 B[...,k] 将 np.view() 返回到 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,5)

如果无法重构代码，则需要开始分析。您可以通过

轻松地临时交换轴

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

与 max9111 的评论相反，与 np.ascontiguousarray() 相比，这不会给您带来任何好处，因为在这两种情况下都必须复制数据。也就是说，副本是 O(Nx*Ny*k) + 缓冲区分配。直接矩阵向量乘法是O(Nx*Ny)，但是你必须先收集元素，这真的很昂贵。这归结为您的特定架构和具体问题，因此分析是可行的方法。