使用转置优化数组加法和乘法答案

【问题标题】：Optimizing array additions and multiplications with transposes使用转置优化数组加法和乘法
【发布时间】：2022-02-21 01:44:21
【问题描述】：

我想在数组转置的基础上做加法和乘法。

鉴于A 是一个数组。 sum += p(A) * factor(p)

其中p 是一个置换/转置，factor(p) 是一个前置因子列表。比如A是一个二维数组，目标是

0.1*A + 0.1*transpose(A,(1,0))

我在我的Python 代码中发现，当应用更高维排列时，数组加法的时间远大于转置。也许numpy.transpose 使用C 中的指针。我想知道，有没有什么方法可以优化数组加法部分的时序？ numpy.add 没有多大帮助。我应该以某种方式只对影响数组的转置部分求和，其余部分使用乘法吗？例如，排列(0,1,3,2)，(0,1) 部分在前一个数组的顶部乘以一个公因子。或者使用cpython来提升性能？

这是我的Python 代码

import numpy as np
import time
import itertools as it

ref_list = [0, 1, 2, 3, 4]
p = it.permutations(ref_list)
transpose_list = tuple(p)
print(type(transpose_list),type(transpose_list[0]),transpose_list[0])


n_loop = 2
na = nb = nc = nd = ne = 20
A = np.random.random((na,nb,nc,nd,ne))
sum_A = np.zeros((na,nb,nc,nd,ne))
factor_list = [i*0.1 for i in range(120)]

time_transpose = 0
time_add = 0
time_multiply = 0

for n in range(n_loop):
    for m, t in enumerate(transpose_list):
        start = time.time()
        B = np.transpose(A, transpose_list[m] )  
        finish = time.time()
        time_transpose += finish - start

        start = time.time()
        B_p = B * factor_list[m]  
        finish = time.time()
        time_multiply += finish - start

        start = time.time()
        sum_A += B_p
        finish = time.time()
        time_add += finish - start

print(time_transpose, time_multiply, time_add, time_multiply/time_transpose, time_add/time_transpose)

输出是

0.004961967468261719 1.1218750476837158 3.7830252647399902 226.09480107630213 762.404285988852

加法时间比转置大约大 700 倍。

我尝试在How to avoid huge overhead of single-threaded NumPy's transpose?中使用numba的转置

通过添加


import numba as nb

@nb.njit('void(float64[:,::1], float64[:,::1])', parallel=True)
def transpose(mat, out):
    blockSize, tileSize = 256, 32  # To be tuned
    n, m = mat.shape
    assert blockSize % tileSize == 0
    for tmp in nb.prange((m+blockSize-1)//blockSize):
        i = tmp * blockSize
        for j in range(0, n, blockSize):
            tiMin, tiMax = i, min(i+blockSize, m)
            tjMin, tjMax = j, min(j+blockSize, n)
            for ti in range(tiMin, tiMax, tileSize):
                for tj in range(tjMin, tjMax, tileSize):
                    out[ti:ti+tileSize, tj:tj+tileSize] = mat[tj:tj+tileSize, ti:ti+tileSize].T

并使用

B = transpose(A, transpose_list[m] )

收到

Traceback (most recent call last):
  File "transpose_test_v2.py", line 46, in <module>
    B = transpose(A, transpose_list[m] )
  File "/home/.../lib/python3.8/site-packages/numba/core/dispatcher.py", line 717, in _explain_matching_error
    raise TypeError(msg)
TypeError: No matching definition for argument type(s) array(float64, 6d, C), UniTuple(int64 x 6)

或使用 B = nb.transpose(A, transpose_list[m] ) 并得到了

    B = nb.transpose(A, transpose_list[m] )
AttributeError: 'int' object has no attribute 'transpose'

【问题讨论】：

标签： python numpy optimization transpose

【解决方案1】：

np.transpose 不会物理转置内存中的数组。它改变了数组的步幅，以便（实际上）转置结果视图。问题是视图的内存布局在转置后非常低效，因此对其进行的操作效率不高。您可以复制转置视图或致电np.ascontiguousarray 以便急切地执行转置。请注意，Numpy copy of transposed arrays is quite inefficient 的当前实现。

这是我机器上的结果：

Initial code:
0.0030066967010498047 1.2316536903381348 1.971841812133789 409.6368249940528 655.8166679882642

With a copy of the transposed array:
2.483657121658325 1.221949577331543 0.6028566360473633 0.4919960837893974 0.2427294133277298

除此之外，预分配数组并使用np.add和np.multiply的参数out应该有助于加快乘法和加法的速度。为了获得更快的性能，您可以使用带有循环的并行 Numba 代码。由于维数较多，因此在您的情况下优化转置非常困难。

更新：

这是一个更快的 Numba 实现，它主要消除了加法/乘法的成本，并降低了使用多线程的转置成本：

# Same setup than in the question

@numba.njit('(float64[:,:,:,:,::1], float64[:,:,:,:,::1], float64, int_, int_, int_, int_, int_)', parallel=True, fastmath=True)
def computeRound(A, sum_A, factor, i0, i1, i2, i3, i4):
    size = 20
    B = np.transpose(A, (i0, i1, i2, i3, i4))
    for j0 in numba.prange(size):
        for j1 in range(size):
            for j2 in range(size):
                for j3 in range(size):
                    for j4 in range(size):
                        sum_A[j0, j1, j2, j3, j4] += B[j0, j1, j2, j3, j4] * factor

for n in range(n_loop):
    for m, t in enumerate(transpose_list):
        i0, i1, i2, i3, i4 = transpose_list[m]
        computeRound(A, sum_A, factor_list[m], i0, i1, i2, i3, i4)

此代码比初始代码快 4 倍。它几乎完全受转置的不良访问模式的约束。然而，这很难优化转置。一种解决方案是以从先前的转置最大化缓存局部性的顺序应用转置（快速但特别难以实现）。一个更简单（但速度较慢）的解决方案是使用先前提供的优化 Numba 转置代码，同时手动处理 (i0, i1, i2, i4, i3) 案例而不是 (i0, i1, i2, i3, i4) 案例。

【讨论】：

非常感谢。我在循环上方添加了A = np.ascontiguousarray(A, dtype=np.float32)，在B = np.ascontiguousarray(B, dtype=np.float32) 之后添加了B = np.ascontiguousarray(B, dtype=np.float32)，我得到了和你类似的输出。添加时间减少，但B = np.ascontiguousarray(B, dtype=np.float32) 需要类似（略少）的时间。总时间不会减少太多。有什么更好的方法吗？加载 C/Fortran 是否有助于优化加法/乘法部分？ np.add(sum_A, B_p, out = sum_A ) 没有帮助:(
这是正常的ascontiguousarray 并没有快多少，因为这不是使用它的目的。关键是要表明由于换位的工作方式，您的测量是有偏差的。答案的第二部分指出了一些获得更快结果的方法。至于转置，我指出的 Numba 解决方案（请参阅答案的链接）应该使其更快。请注意，最近版本的 Numba 优化了转置数组的副本，而不是 Numpy，因此它应该更快地使用它，而且只需一点点努力。尽管如此，转置是非常昂贵的操作（如矩阵乘法）。
非常感谢。我尝试使用numba 并没有得到它的工作:( 在操作中编辑。
这是完全正常的 Numba 代码在您的情况下不起作用。 Numba 代码旨在在您处理 5D 数组时处理 2D 数组（请参阅显式签名和链接问题）。您需要调整代码，而不仅仅是复制过去。由于维度数量众多，这显然不是一件简单的事情。
@Geositta 我使用 Numba 实现了一个明显更快的版本。不是银弹，而是更好。我还提供了优化这一点的技巧。

【解决方案2】：

您的错误消息看起来像“(m+blockSize-1)//blockSize”部分，其中除法为您提供双精度浮点数。不用除以blockSize(256)，你可以右移8次并保持整数。

【讨论】：