NumPy中整数数组索引数组的矩阵乘法性能退化问题

问题背景 #

在项目中需先对数组进行行/列置换，再执行（广播）矩阵乘法。使用NumPy实现时发现，置换方式不同会导致性能差异极大，以下是带计时的最小示例：

import time
import numpy as np

n = 800
nt = 50

a = np.random.randn(10, n)
b = np.random.randn(7, n, n)
p = np.random.permutation(n)

# 基准测试：无置换
t0 = time.time()
c = [a @ b for i in range(nt)]
t1 = time.time()
print('Time: ', t1-t0)          # ~ 0.2 seconds

# 仅置换最后一维
b1 = b[:, :, p]
t0 = time.time()
c = [a @ b1 for i in range(nt)]
t1 = time.time()
print('Time: ', t1-t0)          # ~ 0.22 seconds

# 先置换中间维度，再置换最后一维
b2 = b[:, p, :][:, :, p]
t0 = time.time()
c = [a @ b2 for i in range(nt)]
t1 = time.time()
print('Time: ', t1-t0)          # ~ 4.1 seconds

# 先置换最后一维，再置换中间维度
b3 = b[:, :, p][:, p, :]
t0 = time.time()
c = [a @ b3 for i in range(nt)]
t1 = time.time()
print('Time: ', t1-t0)          # ~ 12.5 seconds

通过查看数组的strides发现，性能差异源于内存布局的变化：

print('b: ', b.__array_interface__['strides'])      # None（表示C连续）
print('b1: ', b1.__array_interface__['strides'])    # (6400, 8, 44800)
print('b2: ', b2.__array_interface__['strides'])    # (8, 56, 44800)
print('b3: ', b3.__array_interface__['strides'])    # (8, 44800, 56)

上述示例中，仅置换3维数组最后一维时性能基本无变化，但同时置换最后两维时性能显著下降，且置换顺序也会影响性能。所有置换后的数组均非C风格连续数组，此情况超出预期。

一、整数数组索引的工作原理 #

NumPy的整数数组索引默认生成**视图（View）**而非副本：它不会拷贝原数组的内存，而是通过修改数组的strides（内存步长）和形状，来映射原数组的元素位置，以此避免内存拷贝的开销。

但这种视图的内存布局可能完全失去连续性：

• 当仅对最后一维做索引（如b[:, :, p]），原数组的行是连续存储的，即使行内元素被打乱，同一行的内存访问局部性仍较好，CPU缓存命中率高，所以性能下降不明显。
• 当对中间维度做索引，或多次嵌套索引时，生成的视图strides会变得极不规则。比如b3的中间维度步长为44800字节，意味着每次访问中间维度的下一个元素时，需要跳跃大量内存，彻底破坏CPU缓存的局部性——而矩阵乘法严重依赖缓存效率，这直接导致性能断崖式下跌。

二、规避性能退化的方法 #

1. 合并多次索引为单次操作 #

把多步置换合并成一次索引，减少视图的内存布局扭曲。例如把b[:, p, :][:, :, p]改为b[:, p, p]，仅生成一次视图：

b2_optimized = b[:, p, p]
t0 = time.time()
c = [a @ b2_optimized for i in range(nt)]
t1 = time.time()
print('Optimized b2 Time: ', t1-t0)  # 性能接近b1的水平

2. 强制转换为连续数组 #

对于已经生成的非连续视图，使用np.ascontiguousarray()或.copy()方法将其转换为C风格连续数组。虽然会产生一次内存拷贝，但后续矩阵乘法的性能会大幅回升：

# 优化b2
b2_contiguous = np.ascontiguousarray(b2)
t0 = time.time()
c = [a @ b2_contiguous for i in range(nt)]
t1 = time.time()
print('Contiguous b2 Time: ', t1-t0)  # 耗时回到0.2秒左右

# 优化b3
b3_contiguous = np.ascontiguousarray(b3)
t0 = time.time()
c = [a @ b3_contiguous for i in range(nt)]
t1 = time.time()
print('Contiguous b3 Time: ', t1-t0)  # 性能显著提升

3. 调整置换顺序，优先保障关键维度连续性 #

矩阵乘法中，参与运算的核心维度（此处为b的最后两维）的内存连续性对性能影响最大。尽量让置换后的数组最后两维保持连续；若无法避免非连续，提前通过拷贝转换为连续数组，再执行乘法操作。

内容的提问来源于stack exchange，提问作者jinzx10