从 0.98 到 267 GFLOPS:矩阵乘法性能优化全记录

Relativity by M.C. Escher
M.C. Escher,《相对性》, 1953

引言

矩阵乘法(GEMM, General Matrix Multiply)是高性能计算中最核心的操作之一。神经网络的全连接层、卷积层、注意力机制,底层都是矩阵乘法。BLAS 库里的 sgemm 是被优化得最极致的函数之一,Intel MKL、OpenBLAS、BLIS 这些库的核心竞争力就在这里。

我想搞清楚,从一个教科书上的三重循环,到一个接近硬件理论峰值的实现,中间到底经历了什么。于是我用 C++ 从零开始,一步一步优化,把每一步的思路、效果、踩过的坑都记录下来。

问题定义

计算 C = A × B,其中 A 是 M×K 矩阵,B 是 K×N 矩阵,C 是 M×N 矩阵。所有矩阵都是行优先(row-major)存储的 float32。

硬件环境:

理论峰值是怎么算的?

这里需要拆开讲。i7-14700KF 的单核睿频频率约为 5.5 GHz,意味着每秒执行 55 亿个时钟周期。关键问题是:每个周期能做多少次浮点运算?

这颗 CPU 支持 AVX2 指令集。AVX2 的寄存器宽度是 256 位(32 字节),能同时装 8 个 float32。它还支持 FMA(Fused Multiply-Add)指令——一条指令完成一次乘法和一次加法,也就是 a * b + c,算 2 次浮点运算。而 Raptor Lake 的每个核心有 2 个 FMA 执行单元,可以在同一个时钟周期内同时执行 2 条 FMA 指令。

所以单核每周期的浮点运算数 = 2 个 FMA 单元 × 8 个 float/FMA × 2 FLOP/FMA = 32 FLOP/cycle。理论单核峰值:

5.5 GHz × 32 FLOP/cycle ≈ 176 GFLOPS (单精度, 单核)

多核理论峰值更高,但受限于内存带宽和共享缓存,实际能达到的性能取决于计算密度。后面我们会看到,单线程优化到 144 GFLOPS(峰值的 82%),已经是非常好的结果。

基准测试矩阵大小为 N=1024,即 1024×1024 的方阵相乘。矩阵乘法的计算量公式是 2×M×N×K(每个输出元素需要 K 次乘法和 K 次加法),所以总计算量为 2×1024³ ≈ 21.5 亿次浮点运算。

总览

版本 时间 (ms) GFLOPS 相对 Baseline
Baseline (i-j-k) 2205 0.98
Naive (i-k-j) 61.0 35.3 36×
Tiled 73.5 29.3 30×
Tiled + SIMD 72.8 29.3 30×
Register Tiling 16.9 125.8 128×
Packed 14.8 144.4 147×
Parallel (OpenMP) 8.24 267.2 273×
Optimized 8.21 266.3 272×

Step 0: 教科书写法 (i-j-k) — 0.98 GFLOPS

所有人学线性代数时写的第一个版本:

void matmul_baseline(const float* A, const float* B, float* C,
                     size_t M, size_t N, size_t K) {
  std::memset(C, 0, M * N * sizeof(float));
  for (size_t i = 0; i < M; ++i) {
    for (size_t j = 0; j < N; ++j) {
      for (size_t k = 0; k < K; ++k) {
        C[i * N + j] += A[i * K + k] * B[k * N + j];
      }
    }
  }
}

0.98 GFLOPS,理论峰值的 1% 左右。慢得令人发指。

问题出在内层循环的内存访问模式。我们的矩阵是按行存储的(row-major),也就是说 B[0][0], B[0][1], B[0][2], ... 这些元素在内存中是连续的。但内层循环扫描 k,每次访问 B[k * N + j]——k 每增加 1,内存地址就跳了 N 个 float(4096 字节),相当于一行一行地跳。

现代 CPU 有一个重要的设计:当你从内存读一个数据时,CPU 不会只读那一个数,而是把它所在的一整条 cache line(通常 64 字节)都加载进来。对于 float32,一条 cache line 能装 16 个连续的 float。CPU 的假设是"你读了这个数,大概率接下来会读它旁边的数"——这叫空间局部性(spatial locality)。

但 i-j-k 循环完全违背了这个假设。每次加载一条 cache line(64 字节 = 16 个 float),我们只用了其中 1 个就跳走了。cache line 利用率只有 1/16 = 6.25%。Cachegrind 实测 D1 miss rate 高达 32.8%——每 3 次内存访问就有 1 次 cache miss。

Step 1: 循环重排 (i-k-j) — 35.3 GFLOPS

k 循环提到中间,让内层循环扫描 j

void matmul_naive(const float* A, const float* B, float* C,
                  size_t M, size_t N, size_t K) {
  std::memset(C, 0, M * N * sizeof(float));
  for (size_t i = 0; i < M; ++i) {
    for (size_t k = 0; k < K; ++k) {
      float a_ik = A[i * K + k];
      for (size_t j = 0; j < N; ++j) {
        C[i * N + j] += a_ik * B[k * N + j];
      }
    }
  }
}

36 倍加速,仅仅靠交换两行循环。

现在内层循环里 B[k*N+j]C[i*N+j] 都是沿着行连续访问的(stride-1)。每条 64 字节的 cache line 加载进来以后,里面 16 个 float 全部用到,cache line 利用率从 6.25% 提升到 100%。同时 A[i*K+k] 在内层循环中是常量,被提到循环外保存在一个局部变量 a_ik 里——这叫 scalar replacement,编译器通常也会自动做这件事。

Cachegrind 实测 D1 miss rate 从 32.8% 降到 4.8%

这是整个项目中单步提升最大的优化。教训很简单:内存访问模式比算法复杂度更重要。两种写法的算术运算量完全一样(都是 2N³ 次),但内存访问模式的差异造成了 36 倍的性能差距。

Step 2: Cache Tiling — 29.3 GFLOPS

教科书上的经典优化:把大矩阵分成小块(64×64×64),让每个块的数据量(64×64×4B = 16KB)能放进 48KB 的 L1d cache。这样在处理一个小块时,数据可以反复使用而不被换出 cache。

void matmul_tiled(const float* A, const float* B, float* C,
                  size_t M, size_t N, size_t K) {
  constexpr size_t BM = 64, BN = 64, BK = 64;
  std::memset(C, 0, M * N * sizeof(float));

  for (size_t i0 = 0; i0 < M; i0 += BM) {
    const size_t imax = std::min(i0 + BM, M);
    for (size_t k0 = 0; k0 < K; k0 += BK) {
      const size_t kmax = std::min(k0 + BK, K);
      for (size_t j0 = 0; j0 < N; j0 += BN) {
        const size_t jmax = std::min(j0 + BN, N);
        for (size_t i = i0; i < imax; ++i) {
          for (size_t k = k0; k < kmax; ++k) {
            float a_ik = A[i * K + k];
            for (size_t j = j0; j < jmax; ++j) {
              C[i * N + j] += a_ik * B[k * N + j];
            }
          }
        }
      }
    }
  }
}

结果:反而变慢了。从 35.3 降到 29.3 GFLOPS。

原因是编译器(GCC/Clang -O3)对简单的 i-k-j 循环已经做了很好的优化。分块引入了额外的循环控制开销和 std::min 边界计算,得不偿失。而且 i-k-j 本身的访问模式已经很友好——B 和 C 都是 stride-1 访问,tiling 的边际收益不大。

教训:Tiling 的价值要配合 register tiling 和 packing 才能显现。单独使用 tiling,现代编译器可能做得比你更好。

Step 3: 手写 SIMD (AVX2 FMA) — 29.3 GFLOPS

先解释一下 SIMD 是什么。SIMD(Single Instruction, Multiple Data)是一种并行计算方式——一条指令同时处理多个数据。普通标量指令一次只能对 1 个 float 做加法;而 AVX2 的 SIMD 指令可以一次对 8 个 float 同时做加法(因为 YMM 寄存器是 256 位宽,256 / 32 = 8 个 float)。

FMA(Fused Multiply-Add)更进一步:一条指令完成 a * b + c,也就是同时做乘法和加法,而且只有一次舍入误差(比先乘后加更精确)。对矩阵乘法来说,核心操作就是 C[i][j] += A[i][k] * B[k][j],天然适合 FMA。

在 C/C++ 中,可以通过 intrinsics 函数直接调用这些指令,而不用写汇编:

// 内层循环核心
__m256 va = _mm256_set1_ps(A[i * K + k]);     // 把 1 个 float 广播到 8 路
size_t j = j0;
for (; j + 8 <= jmax; j += 8) {
  __m256 vb = _mm256_loadu_ps(&B[k * N + j]); // 一次加载 8 个连续 float
  __m256 vc = _mm256_loadu_ps(&C[i * N + j]); // 一次加载 8 个连续 float
  vc = _mm256_fmadd_ps(va, vb, vc);            // 8 路并行: C += A * B
  _mm256_storeu_ps(&C[i * N + j], vc);         // 一次写回 8 个 float
}

这里每个 _mm256 开头的函数都对应一条 AVX2 指令。_mm256_set1_ps 把一个标量广播到 256 位寄存器的所有 8 个位置;_mm256_loadu_ps 从内存加载 8 个连续 float 到寄存器;_mm256_fmadd_ps 对 8 对数据同时执行 FMA。

结果:和 tiled 版本完全一样

-O3 -march=native -ffast-math 下,编译器已经自动向量化(auto-vectorization)了内层循环。也就是说,编译器看到了这个简单的 for j 循环,自动把它转换成了 SIMD 指令——和我们手写的 intrinsics 做的事情一模一样。可以用 -fopt-info-vec(GCC)或 -Rpass=loop-vectorize(Clang)来确认编译器是否做了自动向量化。

教训:在现代编译器面前,简单循环的手写 SIMD 通常无法超越自动向量化。真正的收益来自编译器做不到的事——比如下一步的 register tiling,它需要程序员精确规划每个寄存器的用途。

Step 4: Register Tiling — 125.8 GFLOPS(关键突破)

前面所有版本都有一个根本问题:每做一次乘加,都要从内存(或 cache)读写 C。在 i-k-j 的内层循环中,C[i*N+j] += a_ik * B[k*N+j] 这行代码意味着每次都要读 C、做乘加、写回 C。即使 C 在 L1 cache 中,读写延迟也是存在的。

Register tiling 的思路:把 C 的一个小子块完全放在 CPU 寄存器里。寄存器是 CPU 中最快的存储——访问延迟为 0,不经过任何缓存层级。如果我们能让 C 的子块在整个 k 循环中都留在寄存器里,只在循环开始时加载一次、结束时存回一次,就能大幅减少内存访问。

关键问题是:子块能做多大?这取决于 CPU 有多少寄存器。

              NR=16 (2 × __m256)
           +----------------------+
     MR=6  | c00 c01              |  12 个 YMM 寄存器保存 C 的 6×16 子块
    rows   | c10 c11              |
     of C  | c20 c21              |  k 循环中:
           | c30 c31              |    载入 2 个 B 向量 (b0, b1)
           | c40 c41              |    广播 6 个 A 标量
           | c50 c51              |    12 次 FMA
           +----------------------+

AVX2 提供了 16 个 YMM 寄存器(ymm0-ymm15),每个 256 位宽(8 个 float)。我们需要精打细算地分配它们:

为什么选 MR=6, NR=16 而不是其他组合?因为这个组合的计算密度最高。每步 k 做 12 次 FMA(6 行 × 2 个向量),只需从内存读 3 个值(2 个 B 向量 + 1 个 A 标量),计算密度为 12/3 = 4 FMA/load。如果用 4×24(也是 12 个 C 寄存器),则需要读 3 个 B + 1 个 A = 4 次,密度为 12/4 = 3 FMA/load,更低。

核心实现(micro-kernel):

// 1. 加载 C 的 6×16 子块到 12 个寄存器(只做一次)
__m256 c00 = _mm256_loadu_ps(&C[0 * N + 0]);
__m256 c01 = _mm256_loadu_ps(&C[0 * N + 8]);
__m256 c10 = _mm256_loadu_ps(&C[1 * N + 0]);
__m256 c11 = _mm256_loadu_ps(&C[1 * N + 8]);
// ... c20, c21, c30, c31, c40, c41, c50, c51

// 2. k 循环中完全在寄存器里累加
for (size_t k = 0; k < K; ++k) {
  __m256 b0 = _mm256_loadu_ps(&B[k * N + 0]);  // 加载 B 的 16 个 float
  __m256 b1 = _mm256_loadu_ps(&B[k * N + 8]);

  __m256 a;
  a = _mm256_set1_ps(A[0 * lda + k]);  // 广播 A[0][k]
  c00 = _mm256_fmadd_ps(a, b0, c00);   // c00 += a * b0
  c01 = _mm256_fmadd_ps(a, b1, c01);   // c01 += a * b1

  a = _mm256_set1_ps(A[1 * lda + k]);  // 广播 A[1][k]
  c10 = _mm256_fmadd_ps(a, b0, c10);
  c11 = _mm256_fmadd_ps(a, b1, c11);

  // ... 同理处理第 2-5 行,共 12 次 FMA
}

// 3. k 循环结束,把 C 写回内存(只做一次)
_mm256_storeu_ps(&C[0 * N + 0], c00);
_mm256_storeu_ps(&C[0 * N + 8], c01);
// ...

Cachegrind 实测:D refs(数据访问总数)从 177M 降到 21M(减少 8.4 倍)。因为 C 不再在 k 循环中反复读写——以前每步 k 都要读写 C 的 6×16 = 96 个 float,现在整个 k 循环只在开头和结尾各一次。

4.3 倍加速,从 29 GFLOPS 到 126 GFLOPS。这是除了循环重排之外最重要的单步优化。编译器做不了这种级别的寄存器编排——它不会主动把 C 的一大块保持在寄存器里跨越整个循环。这需要程序员对硬件资源有精确的理解。

Step 5: Goto-style Packing — 144.4 GFLOPS

Register tiling 的 micro-kernel 仍然通过 stride=K 访问 A、stride=N 访问 B。当矩阵很大时,相邻的 micro-kernel 调用之间,A 和 B 的数据在内存中不连续——中间隔着很多不需要的数据,浪费了 cache line 的空间。

解决方案来自 Kazushige Goto 在 2008 年的经典论文(GotoBLAS 的作者):在执行 micro-kernel 之前,把 A 和 B 的数据重新排列(pack)到连续的内存区域中,让 micro-kernel 的每次内存读取都是完全连续的。

原始 A (stride = K):              pack_A 后 (stride = MR):
+------------------+              +---------+
| a00 a01 a02 ...  |  MR=6 rows  | a00     |
| a10 a11 a12 ...  |  -------->  | a10     |  每 MR 个值连续
| ...              |              | a20     |  微内核顺序读取
| a50 a51 a52 ...  |              | ...     |
+------------------+              | a01     |
                                  | a11     |
                                  +---------+

但仅仅 packing 还不够,还需要精心设计循环结构,让每层缓存都被充分利用。这就是 Goto-style 的三层循环嵌套:

for (size_t jc = 0; jc < N; jc += NC) {       // 第 3 层: B panel 驻留 L3
  for (size_t kc = 0; kc < K; kc += KC) {     // 第 2 层: B slice 驻留 L2
    pack_B(&B[kc * N + jc], packed_B, kk, nc, N);

    for (size_t ic = 0; ic < M; ic += MC) {   // 第 1 层: A slice 驻留 L1
      pack_A(&A[ic * K + kc], packed_A, mc, kk, K);

      // Macro-kernel: iterate over MR×NR micro-tiles.
      for (size_t ir = 0; ir < mc; ir += MR) {
        for (size_t jr = 0; jr < nc; jr += NR) {
          microkernel_packed_6x16(
              &packed_A[ir * kk],
              &packed_B[jr * kk],
              &C[(ic + ir) * N + (jc + jr)],
              kk, N);
        }
      }
    }
  }
}

每层 tile 的尺寸精确对应一级缓存:

设计的核心思想:packed_B 在最外层循环被打包一次,然后被所有 MC-strip 共享重用。packed_A 在每个 MC-strip 被打包一次,然后在内层循环中被所有 NR-strip 共享重用。这样 packing 的开销被大量的 micro-kernel 计算摊销掉了。

单线程从 126 到 144 GFLOPS,15% 提升。大矩阵上提升更显著,因为大矩阵的数据在 cache 中不连续的问题更严重。

Step 6: 多线程并行 (OpenMP) — 267.2 GFLOPS

前面所有优化都是单线程的。i7-14700KF 有 14 核 20 线程——我们只用了其中一个核心。是时候把其他核心也用起来了。

OpenMP 是 C/C++/Fortran 中最常用的共享内存并行编程接口。它通过编译器指令(pragma)来控制并行化,而不是手动管理线程。只需在循环前加一行 #pragma omp parallel for,编译器就会自动把循环迭代分配到多个线程上执行。

#pragma omp parallel
{
  // 每线程私有的 packed_A buffer
  float* packed_A = static_cast<float*>(
      std::aligned_alloc(32, MC * KC * sizeof(float)));

  #pragma omp for schedule(dynamic)
  for (size_t ic = 0; ic < M; ic += MC) {
    const size_t mc = std::min(MC, M - ic);
    pack_A(&A[ic * K + kc], packed_A, mc, kk, K);

    for (size_t ir = 0; ir < mc; ir += MR) {
      for (size_t jr = 0; jr < nc; jr += NR) {
        microkernel_packed_6x16(
            &packed_A[ir * kk],
            &packed_B[jr * kk],
            &C[(ic + ir) * N + (jc + jr)],
            kk, N);
      }
    }
  }

  std::free(packed_A);
}

并行化的关键设计决策:

从 144 到 267 GFLOPS,1.85 倍加速。没有达到理想的线性加速(14 核应该 14 倍),原因包括:内存带宽成为瓶颈、L3 cache 竞争、线程调度开销,以及 pack_B 仍然是单线程串行执行的。但在 N=1024 这个尺度上,267 GFLOPS 已经是相当好的结果了。

Step 7: 微优化 (Unroll + Prefetch + Huge Pages) — 266.3 GFLOPS

最后尝试三项微优化:

1) k-unroll by 4

把 micro-kernel 的 k 循环展开 4 次。不展开时,每次迭代结束都要执行比较、自增、跳转指令;展开后每 4 次迭代只需要一组控制指令,减少了循环开销。更重要的是,展开给了 CPU 更多指令可以同时执行——现代 CPU 是乱序执行的,它会在一条指令等待数据时先执行后面不依赖它的指令,展开暴露了更多这样的并行机会。

for (; k + 4 <= kc; k += 4) {
  _mm_prefetch(packed_A + 4 * MR, _MM_HINT_T0);
  _mm_prefetch(packed_B + 4 * NR, _MM_HINT_T0);

  MICRO_ITER(packed_A + 0 * MR, packed_B + 0 * NR);
  MICRO_ITER(packed_A + 1 * MR, packed_B + 1 * NR);
  MICRO_ITER(packed_A + 2 * MR, packed_B + 2 * NR);
  MICRO_ITER(packed_A + 3 * MR, packed_B + 3 * NR);

  packed_A += 4 * MR;
  packed_B += 4 * NR;
}

2) Software prefetch

CPU 有硬件预取器(hardware prefetcher),能自动检测连续的内存访问模式并提前加载数据。但硬件预取器不是万能的——它只能识别简单的线性模式(比如连续递增),对于跳跃式访问或分支后的访问模式,它可能预测不准。

Software prefetch 通过 _mm_prefetch 指令,显式告诉 CPU:"我过几步就要用这个地址的数据了,你现在就开始把它从内存加载到 cache 吧。" 这样等到代码真正需要这个数据时,它已经在 cache 里了,避免了等待内存的延迟(通常 50-100ns)。

_MM_HINT_T0 表示加载到 L1 cache(最近、最快),_MM_HINT_T1 表示加载到 L2 cache。我们对下一个 B panel 用 _MM_HINT_T1,因为它比较大,放在 L2 更合适。

3) Transparent Huge Pages

这涉及到虚拟内存的一个关键组件:TLB(Translation Lookaside Buffer,页表缓存)。

现代操作系统使用虚拟内存——程序看到的内存地址不是真实的物理地址,需要通过页表(page table)翻译。这个翻译过程很慢(可能需要访问多级页表),所以 CPU 用 TLB 来缓存最近用过的地址翻译结果。TLB 的容量很小,典型的 L1 dTLB 只有 64-128 个条目。

默认的内存页大小是 4KB。如果 TLB 有 128 个条目,那它能覆盖的内存范围只有 128 × 4KB = 512KB。而我们的 packed_B 有 4MB——远超 TLB 的覆盖范围。当 micro-kernel 遍历 packed_B 时,大量的 TLB miss 会导致频繁的页表查询。

Huge pages 把页大小从 4KB 增大到 2MB。同样 128 个 TLB 条目可以覆盖 128 × 2MB = 256MB,轻松装下 4MB 的 packed_B。

// packed_B (~4MB) 使用 2MB huge pages, 减少 TLB miss
void* ptr = std::aligned_alloc(2 * 1024 * 1024, aligned_bytes);
madvise(ptr, aligned_bytes, MADV_HUGEPAGE);

madvise 是 Linux 系统调用,MADV_HUGEPAGE 告诉内核这段内存希望使用透明大页。注意 packed_A(每线程 72KB)太小了,使用 2MB 大页反而浪费,所以只对 packed_B 使用。

实际效果:在 N=1024 上 267 → 266 GFLOPS,几乎没有提升。在这个尺度上,硬件预取器本身已经做得不错了,packed 后的数据访问模式也很规则。这些微优化在更大矩阵(N=2048+)上收益更显著——数据量越大,TLB miss 和 prefetch 的影响越明显。

Cachegrind 缓存分析

Valgrind 是一个程序分析工具框架,其中的 cachegrind 模块专门用于分析程序的缓存行为。它通过模拟 CPU 的缓存层级,统计每一次内存访问是否命中了 L1/L2/L3 cache。

用法很简单:

# 先编译出优化版本的可执行文件
bazel build -c opt //bench:cache_profile

# 在 cachegrind 下运行,指定要分析的版本和矩阵大小
valgrind --tool=cachegrind ./bazel-bin/bench/cache_profile naive 512

cachegrind 会输出详细的缓存统计。这里的关键指标:

注意 cachegrind 会让程序运行慢 10-50 倍(因为它在模拟缓存),所以用 N=512 而不是 N=1024 来测:

版本 D refs D1 misses D1 miss rate
Baseline (i-j-k) 177M 58.1M 32.8%
Naive (i-k-j) 177M 8.5M 4.8%
Tiled 179M 1.6M 0.9%
Register Tiling 21M 1.6M 7.8%

这张表非常值得仔细看,因为它精确解释了每一步优化的机制:

总结与思考

这个项目最让我印象深刻的几个点:

1. 内存访问模式是第一性能杀手。 仅靠循环重排就获得 36 倍加速。在现代硬件上,缓存友好的访问模式比任何高级优化都重要。两种写法的算术运算量完全相同,性能差了 36 倍。

2. 不是所有"经典优化"都有效。 Cache tiling 和手写 SIMD 在这个场景下都没有提升。现代编译器已经非常强了,不要低估 -O3 -march=native -ffast-math 的威力。做优化之前先看看编译器已经做了什么。

3. Register tiling 是关键突破。 这是编译器做不到的优化——需要程序员精确理解硬件有多少寄存器、每个寄存器该放什么、计算密度怎么最大化。12 个 YMM 寄存器保持 C 的子块,消除了 k 循环中对 C 的反复读写,D refs 降了 8.5 倍。

4. Packing 解决的是数据布局问题。 把 A 和 B 重新排列成 micro-kernel 友好的连续内存,让 cache line 的每个字节都被充分利用。Goto-style 的三层循环精确对应 L1/L2/L3 三级缓存,是工程和理论结合的典范。

5. 微优化的边际收益递减。 Unroll、prefetch、huge pages 这些优化在 N=1024 上几乎没有效果。性能优化遵循 Pareto 原则——80% 的收益来自 20% 的关键优化。

从 0.98 到 267 GFLOPS,273 倍的加速。这个过程让我真正理解了 Knuth 那句名言的前半句(人们往往只记得后半句):

"We should forget about small efficiencies, say about 97% of the time: premature optimization is the root of all evil. Yet we should not pass up our opportunities in that critical 3%."

关键是找到那 3%。