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

Answer 1

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

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

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

次佳优化是重构，让batch维度为数组的第一维度。这可以在上述矢量化之上完成，通常是可取的。它看起来像

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)，但是你必须先把元素聚集起来，这真的很昂贵。它归结为您的特定体系结构和具体问题，因此分析是可行的方法。

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

How to make two arrays contiguous so that Numba can speed up np.dot()

python

performance

numpy

dot-product

numba