先亮个观点:矩阵转置这种“一眼看穿”的代码,恰恰是最值得做性能优化的对象之一。别急着反驳,我手头的某图像处理管线里就有一个反复出现的二维浮点矩阵转置环节,最初用教科书式的三层循环直接实现,2048×2048 的 float 矩阵转一次要吃掉接近 180 毫秒;优化到最后,同样规模的转置稳定在 1.1 毫秒左右,整体加速比接近 160 倍。这轮从 180 毫秒到 1.1 毫秒的完整过程,涵盖了访存模式分析、分块、寄存器分块、SIMD、多线程和 non-temporal store 这几套常用手段,适合正在跟数值计算、图像处理、按列统计这类访存密集型代码较劲的 C++ 开发者参考。
很多人觉得转置就是“读一行写一列”的活,算法复杂度都摆在那,能优化到哪里去?实际工程项目里完全不是这么回事。你面对的不只是一个三层循环,而是缓存行、TLB、写分配、指令吞吐、线程调度这些底层资源在互相打架。把这层窗户纸捅破之后,你会发现转置性能差不是数学问题,是访存模式问题。
1. 朴素三层循环的180毫秒到底浪费在哪
1.1 测试环境与基准实现
先交代测试环境,后面所有数据都在这套基准上比较。我用的是一台普通工作站,处理器为 8 核 16 线程、支持 AVX2 的桌面级型号,L1 32KB、L2 256KB/core、L3 16MB,内存带宽大约 25GB/s;编译器是某开源 C++ 编译器,统一使用-O3 -march=native参数。测试矩阵是 2048×2048 的 float,单个矩阵占 16MB,刚好超过 L3 容量,又不至于大到被内存带宽彻底淹没,是个很典型的考察对象。
最朴素的实现长这样:
void transpose_naive(const float* src, float* dst, int n) { for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { dst[j * n + i] = src[i * n + j]; } } }注意我把j放在内层,也就是按行读源矩阵、按列写目标矩阵。这段代码在 O3 下的实测成绩是 180 毫秒左右。换算一下,16MB 源数据加 16MB 目标数据一共 32MB,180 毫秒意味着有效带宽只有不到 0.2GB/s,远低于这台机器的实际内存带宽。问题显然不在计算量,因为浮点加法都没有,每次循环就是一次读加一次写。
1.2 列方向写入击穿缓存行的完整解释
现代 CPU 的缓存行一般是 64 字节,一个 float 占 4 字节,一个缓存行能装 16 个 float。朴素版内层循环访问目标矩阵时,地址是j * n + i,其中j固定、i递增,也就是说循环会在同一列的不同行之间来回跳,每次写入地址相隔n * 4字节。对 2048 来说,这个跨度是 8192 字节,约等于 128 个缓存行。
问题出现在这里:当你写dst[j * n + i]这个元素时,CPU 不可能只把 4 字节单独送到内存,它的最小读写单位是缓存行。普通 store 指令还会触发 write-allocate 机制,也就是先把目标地址所在的整个缓存行从内存读进缓存,再在里面修改那 4 字节。这一行缓存里只写了 4 字节有效数据,另外 60 字节完全是被迫拉进来的。
更致命的是,由于循环跨越的是整行距离,下一次写入又会跳到另一个缓存行,刚才那个只写了 4 字节的缓存行再也没有被访问的机会。于是每次写元素都相当于“拉进一行、写四个字节、然后扔掉”,缓存行的有效利用率只有 1/16。16MB 的目标矩阵实际给内存造成的写分配流量接近 256MB,再加上源矩阵读取,总访存量直接被放大了十几倍。
1.3 写分配和TLB的连锁反应
除了缓存行利用率崩掉,TLB(页表缓存)也在同一时间跟着遭殃。系统默认内存页大小通常是 4KB,2048 个 float 正好 8KB,源矩阵一行跨两个页,目标矩阵一列也跨两个页。内层循环每换一次行,访问地址就要跳到不同的页,TLB 里存的那点页表映射根本不够用,频繁 miss 带来的惩罚比缓存 miss 更严重。
这两件事叠加起来,就解释了一个看起来像连续内存搬运的操作为什么能慢到 180 毫秒。这时候哪怕编译器把循环优化出花来,只要访存模式本身是畸形的,性能瓶颈依然死死卡在内存子系统。所以第一步优化不是在指令层面抠细节,而是重新设计访存顺序。
2. 第一波优化:循环分块让数据待在缓存里
2.1 为什么分块有效:两个矩阵块的三份副本
分块(blocking/tiling)的核心思想是让同一块数据在进入缓存之后,被足够多的读写操作复用,而不是用一次就被冲掉。对于转置来说,源矩阵的读是行连续的,目标矩阵的写也是行连续的,单独看都挺友好;但合在一起,源和目标的地址方向正好垂直,整个矩阵太大时彼此会互相驱逐。
解决办法是把大矩阵切成 B×B 的小块。处理一个块时,源块大小占B*B*4字节,目标块同样占B*B*4字节,加上 write-allocate 造成的写缓冲副本,实际上需要在缓存里同时容纳约三份块数据。所以块大小必须满足“三份副本能塞进 L2 缓存”这个条件。
void transpose_blocked(const float* src, float* dst, int n, int block) { for (int i0 = 0; i0 < n; i0 += block) { for (int j0 = 0; j0 < n; j0 += block) { for (int i = i0; i < i0 + block; ++i) { for (int j = j0; j < j0 + block; ++j) { dst[j * n + i] = src[i * n + j]; } } } } }代码结构没有本质变化,就是把两层内循环限制在一个方块内,让源块和目标块在执行期间一直留在 L2 缓存里。这个阶段完全不碰指令集,纯粹靠改善数据局部性,就能获得非常可观的收益。
2.2 块大小怎么选:经验公式与实际扫描
块大小不能拍脑袋定。经验公式是让2 * B * B * 4不超过 L2 容量的一半。这台机器 L2 是 256KB,那么2 * B² * 4 < 128KB,解出来 B 小于 128。我在项目里实际扫描了 16、32、64、128、256 几档,结果很有意思:
| 块大小 | 2048矩阵耗时 | 相对朴素版加速比 |
|---|---|---|
| 16 | 52 ms | 3.5x |
| 32 | 47 ms | 3.8x |
| 64 | 45 ms | 4.0x |
| 128 | 48 ms | 3.8x |
| 256 | 71 ms | 2.5x |
64 是这台机器上的甜点值。块太大时,两块数据加写副本又超出缓存容量,局部性优势打了折扣;块太小时,块间循环的寻址开销占比上升,预取器也来不及识别规律。这个参数在不同 L2 大小的处理器上需要重新扫一遍,不能照搬。
2.3 分块后的实测:访存模式对了,一切都会跟上来
分块版本把 180 毫秒压到了 45 毫秒,加速比整整 4 倍,而这条优化路径没有引入任何花哨指令。它验证了一个我在很多性能调优案例里反复提到的判断:当算法本身就是 O(n²) 的密集访存时,先治访存模式,再考虑指令集和并行。否则 SIMD 加得再漂亮,数据还在内存里做布朗运动,吞吐照样上不去。
这个阶段还有一个附带收益:分块之后 TLB miss 也大幅减少了。每处理一个 64×64 的块,涉及的页数量从全矩阵扫描时的几千次访问降到一个可控范围,页表缓存的压力自然缓解。所以分块不仅是缓存优化,同时也是 TLB 优化。
3. 第二波优化:寄存器分块让访存指令更少、更整齐
3.1 8x8小块在CPU内部完成转置
分块解决了缓存层面的问题,但内层仍然是一个元素一个元素地读和写。每个元素都要生成一条 load 和一条 store 指令,4 字节的访存粒度在 CPU 看来依然很浪费。下一步是在每个 64×64 块内部再做微分块,一次处理 8×8 的浮点数据。
之所以选 8×8,是因为一块数据正好 256 字节,读起来是 4 个连续的缓存行,目标写入也是 4 个连续缓存行。处理过程中这 64 个 float 如果被编译器放进寄存器或栈顶的临时数组,就可以避免在内存里反复倒腾。
void transpose_block_8x8(const float* src, float* dst, int n, int i0, int j0) { float tmp[8][8]; for (int i = 0; i < 8; ++i) { for (int j = 0; j < 8; ++j) { tmp[j][i] = src[(i0 + i) * n + (j0 + j)]; } } for (int i = 0; i < 8; ++i) { for (int j = 0; j < 8; ++j) { dst[(j0 + i) * n + (i0 + j)] = tmp[i][j]; } } }注意第一段循环里tmp[j][i] = src[...]的下标顺序:我读进来的是第i行的连续 8 个元素,但存进临时数组时把行和列对调了。这样第二段循环向目标矩阵写数据时,写入的也是连续 8 个元素。读连续、写连续,中间只隔一个高速暂存区,这是这段代码的精髓。
3.2 为什么不让编译器自动做这件事
很多读者会问:现代编译器不是有自动向量化和循环展开吗,为什么要手工写 8×8?实测下来,编译器在简单两层循环里确实可以展开一部分,但转置的数据流是“行读列写”,列索引的跳跃模式天然不满足自动向量化的连续性检查。即使开了-O3 -march=native,编译器往往只能做部分展开,生成不了漂亮的连续访存序列。
手工写 8×8 块的另一个好处是给编译器一个明确的形状提示:读取阶段是 8 个连续行片段,写入阶段也是 8 个连续行片段,编译器可以更容易地把临时数组的元素调度到寄存器里。我试过 4×4、16×8、8×16 几种形态,8×8 在编译器生成的代码质量和缓存行利用之间达到了最好的平衡。
3.3 实测收益:从45毫秒到20毫秒的原因
寄存器分块后的成绩是 20 毫秒左右,比纯分块又提升了 2 倍多。这部分收益来自两个方面:一是访存指令数量下降,一次处理 64 个元素而不是一个元素一条 load 一条 store,指令前端带宽的压力小了很多;二是短促连续的访存序列更容易被硬件预取器识别,内存控制器可以把请求合并成更高效的 burst。
走到这一步时,访问总量已经接近“读一遍、写一遍”的理论下限,主要的浪费只剩 write-allocate 带来的写前读和缓存行部分利用。下一步的优化空间就轮到 SIMD 和存储指令本身了。
4. 第三波优化:SIMD转置,用shuffle和unpack完成数据重排
4.1 为什么需要两条指令来组合一列
SIMD 化的核心是把“内存搬运”变成“寄存器内数据重排”。传统思路上,转置需要把数据从内存读进寄存器,再以另一种顺序写回内存;而 SIMD 寄存器本身有 shuffle(同 lane 内换位)、unpack(交叉解包)、permute(跨 lane 换位)等指令,可以一次性完成 4 个、8 个甚至 16 个数据的行列对调。
以 SSE 的 4×4 float 转置为例,一个__m128寄存器能装 4 个 float。第一轮把两行数据做 unpacklo/unpackhi,可以得到两两交叉的中间结果;第二轮再用 shuffle 从两组中间结果里抽出正确的列组合。整个过程不经过内存,直接在寄存器里完成了转置。
#include <xmmintrin.h> void transpose4x4_sse(const float* in, int ldIn, float* out, int ldOut) { __m128 r0 = _mm_loadu_ps(in); __m128 r1 = _mm_loadu_ps(in + ldIn); __m128 r2 = _mm_loadu_ps(in + 2 * ldIn); __m128 r3 = _mm_loadu_ps(in + 3 * ldIn); __m128 t0 = _mm_unpacklo_ps(r0, r1); __m128 t1 = _mm_unpackhi_ps(r0, r1); __m128 t2 = _mm_unpacklo_ps(r2, r3); __m128 t3 = _mm_unpackhi_ps(r2, r3); __m128 s0 = _mm_shuffle_ps(t0, t2, 0x44); __m128 s1 = _mm_shuffle_ps(t0, t2, 0xEE); __m128 s2 = _mm_shuffle_ps(t1, t3, 0x44); __m128 s3 = _mm_shuffle_ps(t1, t3, 0xEE); _mm_storeu_ps(out, s0); _mm_storeu_ps(out + ldOut, s1); _mm_storeu_ps(out + 2 * ldOut, s2); _mm_storeu_ps(out + 3 * ldOut, s3); }掩码0x44和0xEE分别表示“取第一个寄存器的低两个元素加第二个寄存器的低两个元素”和“取第一个寄存器的高两个元素加第二个寄存器的高两个元素”。这四行 shuffle 做完,s0 到 s3 就是转置后的四列。
4.2 8x8转置的组装思路与AVX2的坑
AVX2 比 SSE 复杂一些,因为 256 位寄存器内部其实是两个独立的 128 位 lane,跨 lane 的数据交换需要额外的 permute 指令。我最终没有贴完整的 8×8 AVX2 代码,因为它冗长且平台相关,这里讲清楚组装思路更有价值:先把 8 行数据按低 4 列和高 4 列分成两组,对低组做一轮 4×4 SSE 转置,对高组再做一轮 4×4 转置,然后把两组结果用_mm256_permute2f128_ps交叉拼接成完整 8×8 的输出。
这个坑值得单独提醒:如果直接把 8 个__m256加载进来做 unpack,得到的中间结果在低 lane 和高 lane 各自内部是独立的,不会自动跨 lane 组合出正确列序。第一次写这个内核的人十有八九会在结果矩阵里看到“前半段对了、后半段乱了”的诡异现象。遇到这种问题先想到 lane 边界,不要怀疑算法。
4.3 实测:SIMD阶段收益为何比想象中大
SIMD 版本把 20 毫秒压到 8 毫秒,加速比约 2.5 倍。原本以为访存瓶颈下 SIMD 收益有限,实际却超过“一次处理 4 个 float 所以最多 4 倍”的简单预期。原因是向量化之后,单条 load 指令一次拉 16 字节,store 一次写 16 字节,内存控制器看到的请求数量大幅减少,请求合并效率提高,TLB 和缓存行填充的碎片化也进一步降低。
这个阶段我开始建议把转置内核单独做成一个小函数,用微基准单独测试,而不是直接塞进全矩阵流程里。否则出了性能问题你很难分清是内核的问题,还是外层调度的问题。
5. 第四波优化:多线程与non-temporal store的最后冲刺
5.1 任务划分粒度:把独立的大块交给线程
到 SIMD 阶段,单核计算已经接近访存带宽的上限。想继续压时间,只能让多个核心同时去搬数据。转置的好处是天然可并行:不同位置的块之间没有数据依赖,源矩阵只读,目标矩阵每个线程写自己的区域,不会有竞争。
我用 OpenMP 在分块外层加并行:
#pragma omp parallel for collapse(2) schedule(dynamic, 1) for (int i0 = 0; i0 < n; i0 += BLOCK) { for (int j0 = 0; j0 < n; j0 += BLOCK) { transpose_block_8x8(src, dst, n, i0, j0); } }块大小仍是 64。每个块任务读写的数据量约 32KB,刚好适合 L2 容量,线程间不会互相争抢缓存。schedule(dynamic, 1)是为了应对缓存状态差异造成的耗时抖动:前一个块可能把某个缓存行留在 L1,后一个块没有这个优势,静态调度容易让某些线程收尾时拖长尾巴。
实测从 8 毫秒降到 1.5 毫秒,接近 16 线程的理论加速。这个阶段没有出现伪共享问题,因为不同线程写入的目标块地址相差至少一个缓存行,互不踩踏。
5.2 non-temporal store:绕过write-allocate的最后一击
普通 store 的 write-allocate 机制对转置这种“写完就不再读”的场景是纯粹的浪费。写目标矩阵时,CPU 根本不需要先把旧值读进缓存,因为它马上要被完整覆盖。针对这种情况,x86 提供了 non-temporal store 指令,比如_mm256_stream_ps,它直接把数据写往内存,不占用缓存、不触发写前读。
// 以8x8内核的某一行写回为例 __m256 row = /* 转置后的8个float */; _mm256_stream_ps(&dst[offset], row);使用时有几个硬性条件。第一,目标地址必须对齐到 32 字节,否则行为未定义,我在调试阶段被这个坑炸过好几次;建议分配矩阵时用aligned_alloc(32, ...)。第二,stream 写完之后需要执行_mm_sfence(),保证所有写操作顺序对其他执行单元可见。第三,如果转置结果紧接着要被同一段代码回读,就不能用 stream 写,因为数据不在缓存里,回头读它会直接走内存,反而更慢。
5.3 完整优化链路的实测数据
| 优化阶段 | 主要手段 | 2048×2048耗时 | 累计加速比 |
|---|---|---|---|
| 朴素三层循环 | 无 | 180 ms | 1x |
| 循环分块 | block=64 | 45 ms | 4x |
| 寄存器分块 | 8x8暂存 | 20 ms | 9x |
| SIMD转置内核 | SSE/AVX2 shuffle | 8 ms | 22x |
| 多线程 | OpenMP 16线程 | 1.5 ms | 120x |
| non-temporal store | stream写+sfence | 1.1 ms | 163x |
1.1 毫秒意味着有效访存带宽已经非常接近这台机器 25GB/s 的理论上限:32MB 数据量除以 1.1 毫秒约等于 29GB/s,考虑到非临时写节省了额外读流量,这个数字已经到顶了。换而言之,最后这步已经不是在优化算法,而是在试探硬件的物理边界。
6. 经验与避坑:访存密集优化常见问题盘点
6.1 对齐问题的隐形门槛
整个优化链路里,对齐问题折磨我最久的是 stream store。普通 load/store 有_mm_loadu_ps这种 unaligned 版本兜底,性能损失不大;但_mm256_stream_ps严格要求 32 字节对齐,不满足条件的指针传进去,结果无法预料。分配矩阵时务必要用对齐分配接口,而不能直接new float[]。跨平台时更要注意,某些非 x86 架构对非对齐访存的容忍度远低于 x86。
6.2 测性能时的冷缓存与热缓存陷阱
调优过程中很容易被虚假的“优化效果”带偏。如果先后对比朴素版和优化版,每次都让它们处理同一个矩阵,第二次调用时源数据还躺在缓存里,测试结果会严重偏向后者。正确的做法是每次测试前用一个足够大的缓冲区做缓存冲刷,或者随机打乱多个矩阵的测试顺序,取多次的中位数。我在对比分块参数时就因为热缓存效应误判过 block=16 是最优值,换了冷缓存测试法之后结论完全不同。
6.3 编译器优化报告比猜测靠谱
很多现代编译器提供了向量化诊断信息,比如某个循环是否被向量化、卡在哪个环节。与其对着汇编猜,不如直接看优化报告。矩阵转置这种行读列写的模式,编译器在普通三层循环下基本不会自动生成 shuffle 序列,开了-march=native也一样。如果你发现某些循环被向量化了但反汇编里全是标量指令,大概率就是编译器在假装优化,这时候手工内核才有价值。
6.4 什么时候不值得做这套优化
不是所有转置场景都值得上全套组合拳。矩阵小于 64×64 时,数据都在 L1 缓存里,分块版本就够用了,上多线程反而因为线程创建开销变慢。转置频率极低的场景,比如冷启动时做一次数据预处理,优化到 1.1 毫秒和 20 毫秒对用户体验没有区别,不值得投入精力。另外,如果转置结果马上会被密集读取,non-temporal store 就不适合,因为写后读会让数据重新从内存拉一遍,等于白省了一次写分配。
这轮调优做完,我最大的感受是:性能优化不是堆技巧,而是先画清访存模式,找到数据到底在哪一层缓存被卡住,然后再决定用哪把钥匙。矩阵转置恰好是一道非常标准的训练题,它把缓存、TLB、SIMD、并行和存储指令这些概念全部串在了一起。你在自己的项目里如果遇到某个“看起来很简单却慢得离谱”的算子,不妨先照这条链路走一遍,大概率能少走一轮弯路。