F序NumPy数组赋值性能差异：转置操作为何更快？

问题概述 #

当创建F-contiguous（列优先）数组x = np.array(np.random.rand(2000, 400), order='F')，并生成同布局的掩码y = (x > .5)时，执行x[y] = np.nan的速度明显慢于x.T[y.T] = np.nan：

• 服务器环境（Linux x64、Python 3.6.15、NumPy 1.17.3）：前者耗时约2.45s，后者约1.35s，性能提升近50%；
• 本地环境（Python 3.10.12+NumPy 1.21.5、Python 3.12.3+NumPy 1.26.4）：前者约0.66s，后者约0.53s，性能提升约18%。

perf工具的缓存数据验证了差异根源：

• 服务器环境：非转置操作cache-misses占比62.793%，转置后仅20.939%；
• 本地环境：非转置操作cache-misses占比46.80%，转置后25.50%。

测试代码如下：

import numpy as np

source = np.asfortranarray(np.random.rand(5000, 300))
source[3000][150] = np.nan
y = (source > .5)
y[0][0] = not y[0][0]

def trial1():
    global source, y
    x = np.array(source, order='F', copy=True)
    x[y] = np.nan

def trial2():
    global source, y
    x = np.array(source, order='F', copy=True)
    x.T[y.T] = np.nan

print(f"x flags : {source.flags}")
print(f"y flags : {y.flags}")

import timeit
for i in range(5):
    print(f"Without transpose: {timeit.timeit(trial1, number=100)}")
    print(f"With transpose: {timeit.timeit(trial2, number=100)}")

底层原因分析 #

性能差异的核心是内存访问模式与CPU缓存的匹配度：

1. F序数组的内存布局：F-contiguous数组按列优先存储，即每一列的元素在内存中是连续的，列与列之间则分散在不同的内存块。
2. 非转置操作的内存访问：当执行x[y] = np.nan时，掩码y是F序，选中的元素来自不同列的分散位置。CPU缓存加载的是连续内存块，但这些块中只有少量元素需要修改，导致大量缓存未命中（cache-misses），CPU需要频繁从主存读取数据，耗时剧增。
3. 转置操作的内存访问：NumPy的转置是视图操作，不会复制数据，仅调整数组的strides（步长）和轴顺序。转置后数组变为C-contiguous（行优先），此时掩码y.T选中的元素在内存中是连续的行内元素，缓存加载的连续内存块能覆盖更多需要修改的元素，缓存命中率大幅提升，从而缩短操作时间。
4. 转置的额外开销：转置本身只是修改数组的元数据，几乎没有性能损耗，远低于缓存miss带来的时间浪费。

NumPy源码参考位置 #

要深入理解这一逻辑，可以查看NumPy的以下核心源码文件：

• 掩码赋值逻辑：numpy/core/src/multiarray/item_selection.c中的array_setitem函数，负责处理索引赋值，会根据数组的内存布局（C/F序）和索引的布局优化内存访问路径。
• 转置操作实现：numpy/core/src/multiarray/shape.c中的PyArray_Transpose函数，仅调整数组的轴顺序和步长，不涉及数据复制。
• 数组布局判断：numpy/core/src/multiarray/arrayobject.c中关于数组flags（如F_CONTIGUOUS）的判断逻辑，这是内存访问优化的基础。

优化建议 #

• 对于F序数组的掩码赋值操作，若掩码同是F序，优先转置为C序后操作，充分利用CPU缓存的空间局部性。
• 如果后续操作更适合行优先访问，可以直接创建C-contiguous数组，避免频繁转置。

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