矩阵乘法优化：从三层循环到 Apple Silicon

深度学习里最贵的那几下，多半是矩阵乘法。注意力的 $QK^{\top}$ 和 $PV$、全连接、卷积展开之后，都是

$$C = AB,\qquad C_{ij}=\sum_k A_{ik},B_{kj}$$

算力是 $O(n^3)$，访存是 $O(n^2)$。算术强度大约是 $n/6$（三个矩阵各扫一遍、FP32）。$n$ 小的时候吃带宽，$n$ 大的时候吃算力。优化不是把循环写得更花，而是让每个被请上芯片的数，多干几次活。

这篇文章按「图示原理 → 硬件瓶颈 → 最短代码 → 本机数字」往下走。平台是眼前这台 Mac mini M2 Pro（6 个性能核 + 4 个能效核，16GB 统一内存）。编译器 Apple clang 21，-O3 -ffast-math，方阵 FP32，暖机后取中位数。iOS 和 macOS 共用同一套 Accelerate，后面会单独讲两堵不同的墙。

先看这颗芯片

后面每一步都回扣这里。M2 Pro 上其实有两张地图：数据住在哪一层，以及谁来算。

几个后面会反复用到的数：

• cache line 128B，一行刚好 32 个 FP32
• P 核 L1D 大约 128KB，E 核大约 64KB
• 性能簇 L2 大约 16MB，能效簇大约 4MB，SLC 大约 16MB
• 统一内存带宽大约 200 GB/s
• 计算从下到上：标量 → NEON 128-bit → AMX → GPU / ANE

macOS 和 iOS 的指令集是同源的。差的不是「会不会写 kernel」，是频率墙、热和 QoS。

用 Roofline 把终点立住：naive 贴在带宽墙上，Accelerate 已经顶到 AMX 的算力墙。中间那些优化，就是把点往右上挪。

阶梯是这样的：

不是每一步都会更快。有的步是在给下一步铺路，单独拿出来测会变慢。这本身就是软硬件合在一起时该有的诚实。

1. Naive ijk：按教科书写，按列去跳

图先看。内层循环沿 B 的列往下走：

行主序里，B[k, j] 的下一拍是 B[k+1, j]，中间隔了整行。cache line 请来 32 个数，用完 1 个就丢掉。L1 帮不上忙，预取器也跟不住这种步长。

卡在哪：不是乘法慢，是 B 的跨步访问把带宽浪费了。

1
2
3
4
5
6
7

for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
        float s = 0.f;
        for (int k = 0; k < n; k++)
            s += A[i * n + k] * B[k * n + j];
        C[i * n + j] = s;
    }

本机数字：

$N$ 一大，立刻掉到 2～3 GFLOPS。这就是后面所有加速比的分母。

2. 循环交换 ikj：同一套乘法，换一条路

把 $j$ 收到最里面。固定一个 $A_{ik}$，沿 B 的第 $k$ 行走过去，C 的第 $i$ 行就地累加。

对照更清楚。上面跨行跳，下面顺序扫，算的是同一件事：

卡在哪：还是标量乘法，但 空间局部性回来了。128B 的 cache line 里 32 个 FP32 终于被用上。

1
2
3
4
5
6
7

memset(C, 0, n * n * sizeof(float));
for (int i = 0; i < n; i++)
    for (int k = 0; k < n; k++) {
        float a = A[i * n + k];
        for (int j = 0; j < n; j++)
            C[i * n + j] += a * B[k * n + j];
    }

本机数字（相对 naive）：

只换循环顺序，大矩阵直接快一个数量级。这是整篇文章里，代码改动最小、收益最稳 的一步。后面所有花活，都建立在「先别跟 cache line 作对」之上。

3. Cache 分块：一次只请得起三块

$N=1024$ 时，三个矩阵各 4MB，已经超出单核 L1，也挤得 L2 很满。ikj 靠预取还能维持 29 GFLOPS，但微核如果想把一块 A 反复乘很多块 B，这三块必须先住进缓存。

块大小怎么反推：P 核 L1D 大约 128KB。一块 $128\times 256$ 的 FP32 panel 刚好 128KB，可以钉在 L1；对应的 B panel 放到 L2。A 留下，扫完所有 N 方向的块，再换下一段 K。这就是 Goto / BLIS 那套「外层大块 + 内层微核」的骨架。

卡在哪：工作集超过缓存时，复用半径必须由块大小显式管住。预取器救得了顺序流，救不了「同时想记住三块还在长的矩阵」。

1
2
3
4

for (int i0 = 0; i0 < n; i0 += MC)
    for (int k0 = 0; k0 < n; k0 += KC)
        for (int j0 = 0; j0 < n; j0 += NC)
            /* 对 A[i0:i0+MC, k0:k0+KC] · B[k0:k0+KC, j0:j0+NC] 做 ikj */

本机上，单独给 ikj 套一层 64×64 分块，并没有更快：

这不是分块错了。是 ikj 已经被预取器喂饱，多出来的三重分块循环在收税。分块的价值要等微核出现：只有当你真的会把同一块 A 连乘很多块 B，把它钉在 L1 才划算。下面两步会把这件事做实。

4. Packing：给微核铺一条直道

微核希望每次 load 都是连续的、宽度刚好是 NR。原矩阵里，一个 $K_c\times N_r$ 的 panel 每行只要 8 个数，行与行之间却隔着整行步长。

卡在哪：即便块已经在缓存里，跨步 load 仍会打断微核的流水线，TLB 也更容易 miss。

打包本身是一次 gather。只打包、微核仍是标量，拷贝成本会把收益吃光——本机 packed 标量核只有 7.5 GFLOPS，比 ikj 慢三倍。Packing 不是优化，是 给下一步快微核付的入场费。

5. 寄存器分块：C 的一小窗不要来回写

微核窗口盖在 C 上。4×8 一共 32 个累加器，K 维走完再写回内存。

卡在哪：没有这一层，每次 c[j] += a * b[j] 都要读 C、写 C。C 的流量从 $O(n^2)$ 变成跟着 $K$ 放大。寄存器分块把 C 的写回从「每个 k 一次」收成「每个微核一次」。

本机 4×4 标量寄存器分块：

小矩阵整表住在缓存里，寄存器分块几乎白捡一倍。$N=2048$ 时没有外层 cache 分块护着，窗口外的 A、B 重新开始流，数字掉下去。再一次：微核和分块是一对，拆开都不完整。

6. NEON：一次 FMA 吃 4 个 FP32

Apple Silicon 的 P 核每个周期能发多条 128-bit FMA。标量循环喂不饱它们。一条 NEON 寄存器装 4 个 FP32，vfmaq_f32 做的是「广播一个 $A_{ik}$，乘 B 的一行向量，累加进 C」。

卡在哪：算力墙从「标量 ALU」换成了 向量流水线。要喂饱它，pack 好的 A、B 和寄存器里的 C 必须同时在岗。

1
2
3
4
5
6
7
8
9

float32x4_t c0 = vdupq_n_f32(0), c1 = vdupq_n_f32(0);
for (int k = 0; k < kc; k++) {
    float32x4_t a = vld1q_f32(Ap + k * 4);     /* 4 行 A */
    float32x4_t b0 = vld1q_f32(Bp + k * 8);
    float32x4_t b1 = vld1q_f32(Bp + k * 8 + 4);
    c0 = vfmaq_laneq_f32(c0, b0, a, 0);
    c1 = vfmaq_laneq_f32(c1, b1, a, 0);
    /* 其余 3 行同理 */
}

加上前面的分块和 packing，单线程微核在本机非常稳：

大约 50 GFLOPS，几乎不随 $N$ 抖。说明微核已经把单核 NEON 的数据通路走顺了。6 个 P 核的 NEON 理论峰值大约在 600+ GFLOPS 这个量级，单核 50 离「喂饱一条核」还有距离，但已经不是访存在主导。

7. 多线程：先喂 P 核，并看 QoS

C 按行块切开，每块独立写自己的行，GCD concurrent 往外派。M2 Pro 是 6 个 P 核 + 4 个 E 核。默认队列不一定把活都堆在 P 核上；QOS_CLASS_USER_INTERACTIVE 会明确告诉调度器：这是前台要的算力。

卡在哪：单核 NEON 已经不算访存了，但 6 个 P 核还在闲着。另一侧，线程太细或矩阵太小，派发成本会高过计算。

1
2
3
4
5

dispatch_queue_attr_t attr = dispatch_queue_attr_make_with_qos_class(
    DISPATCH_QUEUE_CONCURRENT, QOS_CLASS_USER_INTERACTIVE, 0);
dispatch_apply((n + MC - 1) / MC, dispatch_queue_create("gemm", attr), ^(size_t bi) {
    /* 对第 bi 个 MC 行块跑 NEON 微核 */
});

本机：

三件事：

1. $N=64$ 时多线程是负优化，派发比计算还贵。
2. 大矩阵大约 5～6×，和 6 个 P 核对得上，E 核没帮上太多忙。
3. 桌面 Mac 上两种 QoS 几乎一样；$N=4096$ 才拉开一点。iOS 后台被降到 E 核时，这条缝会变成一道沟。

8. Accelerate：手写到这里，系统库还要快 50 倍

$N=1024$，手写多线程 NEON 是 280 GFLOPS。同一份数据交给 cblas_sgemm：

2708 GFLOPS，相对 naive 是 1072 倍。

差从哪来？Apple 在性能核旁边放了矩阵协处理器 AMX。它做的不是「一次 4 个数」，而是一列 × 一行，outer product 一次更新整块 C。

1
2

cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
            n, n, n, 1.0f, A, n, B, n, 0.0f, C, n);

本机 Accelerate：

AMX 没有公开指令文档，Accelerate 在矩阵够大时会走它。$N=64$ 就已经 524 GFLOPS，说明阈值很低。$N$ 再大，数字从 2.7 TFLOPS 回到 2 TFLOPS 附近，开始碰到带宽、缓存归属和功耗墙。

手写 NEON 很难追上，不是微核不够认真，是 单元就不一样。产品代码到这一步，默认答案就是调用系统库。

Apple Silicon 上怎么选

通用 GEMM、网络里的一层、要不要上 GPU，不是同一条路。

统一内存不是免费午餐

独立显卡机器上，小 GEMM 先被 PCIe 拷贝淹死。Apple Silicon 省掉了这次拷贝，CPU、GPU、ANE 看见的是同一块统一内存。省的是搬，不是缓存颠簸。Metal 的 MPSMatrixMultiplication 有启动开销，小矩阵往往不如 Accelerate；大矩阵、批量、已经在 GPU 上的图，再值得过去。

ANE 不是通用 GEMM

BNNS 吃 CPU 上的卷积和固定形状的网络层。Core ML 能把部分图卸到 ANE。ANE 喜欢静态形状、量化好的推理，不喜欢你随时改 $M,N,K$ 的通用乘法。动态 GEMM 留在 Accelerate 或 GPU。

SME 是公开的下一站

M4 / A18 上 ARM 的 SME / SME2 把矩阵 tile（ZA）做成了公开指令。思路和 AMX 类似：向量外积填一块累加器。以后手写矩阵核有望不再依赖未文档化的协处理器，但今天这台 M2 Pro 上，公开可用、又真正快的，仍然是 Accelerate。

iOS 和 Mac 是同一套 API，两堵不同的墙

A14 起的 A 系同样有 AMX，调用方式也是 cblas_sgemm / BNNS。不要把桌面 2.7 TFLOPS 写进手机 SLA：几秒热节流、进后台被降 QoS、App 被挂起，都会让数字塌掉。前台、用户正在等的计算，显式用 QOS_CLASS_USER_INTERACTIVE；能接受慢的，别去抢 P 核。

本机总表

平台：Mac mini M2 Pro，FP32，$C=AB$，行主序。单位 GFLOPS。空格是这一档太慢、没跑。

读这张表的方式：

• ikj 是软件侧最值的一刀，大矩阵约 10×。
• 分块 / packing 单独看会变慢，它们是微核的前置条件，不是独立加速器。
• NEON 微核 把单核钉在约 50 GFLOPS。
• GCD 在矩阵够大时吃满 P 核，约 300 GFLOPS。
• Accelerate 再往上一个数量级。那一跳不是循环技巧，是 AMX。

对应回 Roofline：naive 贴着带宽墙；ikj 和 NEON 往右上走；多线程接近 NEON 多核屋顶；Accelerate 换了一块屋顶。

可以带走的几条

1. 先画访问轨迹，再写优化。和 cache line 作对的循环，换顺序比换指令集更值钱。
2. 分块、packing、寄存器窗口、SIMD 是一套，不是四道独立菜。拆开测，有的步会倒退。
3. Apple Silicon 上，教学写到 NEON 微核就够看懂；产品上的通用 GEMM，默认 cblas_sgemm。
4. 小矩阵别开线程，也别急着上 Metal。调用和同步比乘法贵。
5. Mac 上量到的峰值，不能当 iOS 的承诺。同一套 ISA，墙不一样。

公式只需要三个：$C_{ij}=\sum_k A_{ik}B_{kj}$，算术强度大约 $n/6$，以及 $\mathrm{GFLOPS}=2N^3/t$。其余的，都是让数据停在更靠近计算的那一层。

项目链接

文中各阶梯的 C 实现和 CMake bench 在 gmm_opt。