如何避免单线程 NumPy 转置的巨大开销？
How to avoid huge overhead of single-threaded NumPy's transpose?
由于 NumPy 的转置函数，我目前遇到了巨大的开销。我发现这个函数几乎总是 运行 在单线程中，无论转置 matrix/array 有多大。我可能需要避免这种巨大的时间成本。
据我了解，如果 numpy 数组足够大，其他函数如 np.dot 或矢量增量将 运行 并行。某些元素方面的操作似乎在 numexpr 包中并行化得更好，但 numexpr 可能无法处理转置。
我想了解解决问题的更好方法。要详细说明这个问题，
- 有时 NumPy 运行s transpose ultrafast (like
B = A.T) 因为转置张量不用于计算或被转储，并且在这个阶段不需要真正转置数据。调用 B[:] = A.T 时，确实会转置数据。
- 我认为并行转置函数应该是一个解决方案。问题是如何实现。
- 希望解决方案不需要 NumPy 以外的包。 ctype 绑定是可以接受的。希望代码不要太难用也不要太复杂。
- 张量转置是一个加号。虽然转置矩阵的技术也可以用于特定的张量转置问题，但我认为可能很难为张量转置编写一个通用的 API。我实际上还需要处理张量转置，但处理张量会使这个 Whosebug 问题复杂化。
- 如果将来有可能实现并行转置，或者有计划存在？那就不用自己实现转置了 ;)
提前感谢您的任何建议！
当前的解决方法
在我的 Linux 个人计算机上处理模型转置问题（A 的大小约为 763MB），可用 4 核（总共 400% CPU）。
A = np.random.random((9999, 10001))
B = np.random.random((10001, 9999))
D = np.random.random((9999, 10001))
当前的解决方法似乎不够有效。一般来说，如果在 4 核上完全并行化，它应该看到大约 3x~4x 的加速CPU，但我写的代码只获得大约 1.5x。
基线：朴素转置（~255 毫秒）
B[:] = A.T # ~255 ms, time for transpose
D[:] = A[:] # ~ 65 ms, time for coping data
有趣的是，如果A是10000 * 10000方阵，那么转置时间会增加到~310ms。我不知道这里发生了什么。如果矩阵是正方形，甚至 C/ctypes 绑定函数的时间成本也会受到影响（更慢）。
C/ctypes 绑定函数（~145 毫秒）
用下面的OpenMP/BLAS（基本，未优化）写成：
// transpose.c
#include <stdio.h>
#include <cblas.h>
void matrix_transpose(double* A, double* B, size_t d1, size_t d2) {
size_t i;
#pragma omp parallel for shared(A, B, d1, d2) private(i)
for (i = 0; i < d1; ++i) {
cblas_dcopy(d2, A + i*d2, 1, B + i, d1);
}
}
然后执行python代码（4核心线程）
from numpy.ctypeslib import as_ctypes
matrix_transpose = np.ctypeslib.load_library("transpose.so", ".").matrix_transpose
matrix_transpose(
as_ctypes(A), as_ctypes(C),
ctypes.c_size_t(n1), ctypes.c_size_t(n2)) # ~145 ms
使用 C/ctype 绑定可能会有些麻烦，因为它不是纯粹的 python，并且还使用 CBLAS 作为外部包。
多处理（~165 毫秒）
nproc = 4
batch = n1 // nproc + 1 # batch 2500 in this case
slices = [slice(i * batch, min(n1, (i+1) * batch)) for i in range(nproc)]
cB = as_ctypes(B)
pB = sharedctypes.RawArray(cB._type_, cB)
def slice_transpose(s):
B = np.asarray(pB)
B[:, s] = A[s].T
with Pool(nproc) as pool:
pool.map(slice_transpose, slices) # ~165 ms
B = np.asarray(pB)
我猜想对于大型集群 (HPC)，更多 processes/threads 不一定会加速。那么processes/threads的个数怎么设置可能也是个问题
在初始问题后编辑
这个问题不仅与并行化有关，还与缓存感知和基于图块的转置算法有关。可能的解决方案可能是
- 将 numba 代码与基于图块的算法结合使用（来自
的回答）。虽然需要 numpy 以外的包，但它几乎是纯粹的 python.
- 使用优化的 blas 库（例如 MKL 或 cuBLAS，它们实现了自己的矩阵转置 API）和 link 到 python，而不是 CBLAS 或 BLAS。如果要分发此程序，需要准备 makefile 或动态 linked 库。
- 使用
pyscf.lib.transpose (python link, c link) for parallel 3-index tensor transpose M.transpose(0,2,1). I'm somehow a fan of pyscf。它的主要用途是量子或半经验化学计算，但它的一些数值计算实用程序可能为此目的进行了优化。在我在服务器上的测试中，转置 (1502, 601, 601) 张量可能比 MKL mkl_domatcopy（9.19 秒）快两倍（4.09 秒）。
相关算法文章：
相关的 Whosebug 和 github 问题页面：
- How to transpose a matrix in an optimal way using blas?
- NumPy：转置复制性能问题
#18512
高效地计算换位很难。此原语不受计算限制，但 内存限制 。对于存储在 RAM 中的大矩阵（而不是 CPU 缓存）尤其如此。
Numpy 使用基于视图的方法，这种方法在只需要数组的一部分并且非常慢地急切完成计算时非常有用（例如，执行复制时）。在这种情况下执行复制时，Numpy 的实现方式会导致大量 缓存未命中 严重降低性能。
I found this function virtually always run in single-threaded, whatever how large the transposed matrix/array is.
这显然取决于所使用的 Numpy 实现。 AFAIK，一些优化的软件包（如 Intel 的软件包）效率更高，并且更经常利用多线程。
I think a parallelized transpose function should be a resolution. The problem is how to implement it.
是也不是。使用更多线程可能并不需要更快。至少不多，而且不是在所有平台上。 使用的算法远比使用多线程重要。
在现代桌面 x86-64 处理器上，每个内核都可以受其 缓存层次结构 的限制。但是这个限制是相当高的。因此，两个内核通常足以使 RAM 吞吐量接近饱和。例如，在我的（4 核）机器上，顺序副本可以达到 20.4 GiB/s（Numpy 成功达到此限制），而我的（实际）内存吞吐量接近 35 GiB/s。复制 A 需要 72 毫秒，而朴素的 Numpy 转置 A 需要 700 毫秒。即使使用我所有的内核，并行实现也不会快于 175 毫秒，而最佳理论时间为 42 毫秒。实际上，由于缓存未命中和我的 L3 缓存饱和，一个简单的并行实现会比 175 毫秒慢得多。
朴素的转置实现不会 write/read 数据连续。内存访问模式是 strided 并且大多数 cache-lines 被浪费了。正因为如此，数据在巨大的矩阵上从内存中 read/written 多次（在当前使用双精度的 x86-64 平台上通常为 8 次）。 基于图块的转置算法 是防止此问题的有效方法。它还大大减少了缓存未命中。理想情况下，应该使用 分层 切片或 无缓存 Z 切片 模式，尽管这 难以实现 .
这里是一个 基于 Numba 的实现，基于之前的信息：
@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
如果您想要更快的代码，您可以使用非常优化的本机库来实现转置，例如 Intel MKL。此类库通常利用低级处理器特定指令（SIMD 指令和 非临时存储）更有效地使用 caches/RAM。
下面是计时结果（假设输出矩阵已经填入内存）：
Naive Numpy: 700 ms
Above code without Numba (sequential): 287 ms
Numba version (sequential): 157 ms
Numba version (parallel): 104 ms
Very-optimized C code (parallel): 68 ms
Theoretical optimum: 42 ms
由于 NumPy 的转置函数，我目前遇到了巨大的开销。我发现这个函数几乎总是 运行 在单线程中，无论转置 matrix/array 有多大。我可能需要避免这种巨大的时间成本。
据我了解，如果 numpy 数组足够大，其他函数如 np.dot 或矢量增量将 运行 并行。某些元素方面的操作似乎在 numexpr 包中并行化得更好，但 numexpr 可能无法处理转置。
我想了解解决问题的更好方法。要详细说明这个问题，
- 有时 NumPy 运行s transpose ultrafast (like
B = A.T) 因为转置张量不用于计算或被转储，并且在这个阶段不需要真正转置数据。调用B[:] = A.T时，确实会转置数据。 - 我认为并行转置函数应该是一个解决方案。问题是如何实现。
- 希望解决方案不需要 NumPy 以外的包。 ctype 绑定是可以接受的。希望代码不要太难用也不要太复杂。
- 张量转置是一个加号。虽然转置矩阵的技术也可以用于特定的张量转置问题，但我认为可能很难为张量转置编写一个通用的 API。我实际上还需要处理张量转置，但处理张量会使这个 Whosebug 问题复杂化。
- 如果将来有可能实现并行转置，或者有计划存在？那就不用自己实现转置了 ;)
提前感谢您的任何建议！
当前的解决方法
在我的 Linux 个人计算机上处理模型转置问题（A 的大小约为 763MB），可用 4 核（总共 400% CPU）。
A = np.random.random((9999, 10001))
B = np.random.random((10001, 9999))
D = np.random.random((9999, 10001))
当前的解决方法似乎不够有效。一般来说，如果在 4 核上完全并行化，它应该看到大约 3x~4x 的加速CPU，但我写的代码只获得大约 1.5x。
基线：朴素转置（~255 毫秒）
B[:] = A.T # ~255 ms, time for transpose
D[:] = A[:] # ~ 65 ms, time for coping data
有趣的是，如果A是10000 * 10000方阵，那么转置时间会增加到~310ms。我不知道这里发生了什么。如果矩阵是正方形，甚至 C/ctypes 绑定函数的时间成本也会受到影响（更慢）。
C/ctypes 绑定函数（~145 毫秒）
用下面的OpenMP/BLAS（基本，未优化）写成：
// transpose.c
#include <stdio.h>
#include <cblas.h>
void matrix_transpose(double* A, double* B, size_t d1, size_t d2) {
size_t i;
#pragma omp parallel for shared(A, B, d1, d2) private(i)
for (i = 0; i < d1; ++i) {
cblas_dcopy(d2, A + i*d2, 1, B + i, d1);
}
}
然后执行python代码（4核心线程）
from numpy.ctypeslib import as_ctypes
matrix_transpose = np.ctypeslib.load_library("transpose.so", ".").matrix_transpose
matrix_transpose(
as_ctypes(A), as_ctypes(C),
ctypes.c_size_t(n1), ctypes.c_size_t(n2)) # ~145 ms
使用 C/ctype 绑定可能会有些麻烦，因为它不是纯粹的 python，并且还使用 CBLAS 作为外部包。
多处理（~165 毫秒）
nproc = 4
batch = n1 // nproc + 1 # batch 2500 in this case
slices = [slice(i * batch, min(n1, (i+1) * batch)) for i in range(nproc)]
cB = as_ctypes(B)
pB = sharedctypes.RawArray(cB._type_, cB)
def slice_transpose(s):
B = np.asarray(pB)
B[:, s] = A[s].T
with Pool(nproc) as pool:
pool.map(slice_transpose, slices) # ~165 ms
B = np.asarray(pB)
我猜想对于大型集群 (HPC)，更多 processes/threads 不一定会加速。那么processes/threads的个数怎么设置可能也是个问题
在初始问题后编辑
这个问题不仅与并行化有关，还与缓存感知和基于图块的转置算法有关。可能的解决方案可能是
- 将 numba 代码与基于图块的算法结合使用（来自
- 使用优化的 blas 库（例如 MKL 或 cuBLAS，它们实现了自己的矩阵转置 API）和 link 到 python，而不是 CBLAS 或 BLAS。如果要分发此程序，需要准备 makefile 或动态 linked 库。
- 使用
pyscf.lib.transpose(python link, c link) for parallel 3-index tensor transposeM.transpose(0,2,1). I'm somehow a fan of pyscf。它的主要用途是量子或半经验化学计算，但它的一些数值计算实用程序可能为此目的进行了优化。在我在服务器上的测试中，转置 (1502, 601, 601) 张量可能比 MKLmkl_domatcopy（9.19 秒）快两倍（4.09 秒）。
相关算法文章：
相关的 Whosebug 和 github 问题页面：
- How to transpose a matrix in an optimal way using blas?
- NumPy：转置复制性能问题 #18512
高效地计算换位很难。此原语不受计算限制，但 内存限制 。对于存储在 RAM 中的大矩阵（而不是 CPU 缓存）尤其如此。
Numpy 使用基于视图的方法，这种方法在只需要数组的一部分并且非常慢地急切完成计算时非常有用（例如，执行复制时）。在这种情况下执行复制时，Numpy 的实现方式会导致大量 缓存未命中 严重降低性能。
I found this function virtually always run in single-threaded, whatever how large the transposed matrix/array is.
这显然取决于所使用的 Numpy 实现。 AFAIK，一些优化的软件包（如 Intel 的软件包）效率更高，并且更经常利用多线程。
I think a parallelized transpose function should be a resolution. The problem is how to implement it.
是也不是。使用更多线程可能并不需要更快。至少不多，而且不是在所有平台上。 使用的算法远比使用多线程重要。
在现代桌面 x86-64 处理器上，每个内核都可以受其 缓存层次结构 的限制。但是这个限制是相当高的。因此，两个内核通常足以使 RAM 吞吐量接近饱和。例如，在我的（4 核）机器上，顺序副本可以达到 20.4 GiB/s（Numpy 成功达到此限制），而我的（实际）内存吞吐量接近 35 GiB/s。复制 A 需要 72 毫秒，而朴素的 Numpy 转置 A 需要 700 毫秒。即使使用我所有的内核，并行实现也不会快于 175 毫秒，而最佳理论时间为 42 毫秒。实际上，由于缓存未命中和我的 L3 缓存饱和，一个简单的并行实现会比 175 毫秒慢得多。
朴素的转置实现不会 write/read 数据连续。内存访问模式是 strided 并且大多数 cache-lines 被浪费了。正因为如此，数据在巨大的矩阵上从内存中 read/written 多次（在当前使用双精度的 x86-64 平台上通常为 8 次）。 基于图块的转置算法 是防止此问题的有效方法。它还大大减少了缓存未命中。理想情况下，应该使用 分层 切片或 无缓存 Z 切片 模式，尽管这 难以实现 .
这里是一个 基于 Numba 的实现，基于之前的信息：
@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
如果您想要更快的代码，您可以使用非常优化的本机库来实现转置，例如 Intel MKL。此类库通常利用低级处理器特定指令（SIMD 指令和 非临时存储）更有效地使用 caches/RAM。
下面是计时结果（假设输出矩阵已经填入内存）：
Naive Numpy: 700 ms
Above code without Numba (sequential): 287 ms
Numba version (sequential): 157 ms
Numba version (parallel): 104 ms
Very-optimized C code (parallel): 68 ms
Theoretical optimum: 42 ms