从 Gflop/s 到 Tflop/s：在 Swift 中优化矩阵乘法
训练大型语言模型（LLM）本质上是一项巨大的矩阵乘法运算。其核心是对 z += x * y 进行数万亿次的重复循环。对于在 Apple Silicon 上开发的人员来说，挑战常常在于如何在 Swift 的高级安全性与这些工作负载所需的原始性能之间取得平衡。
在最近的一次探索中，开发者 zdw 着手将 Andrej Karpathy 的 llm.c（一个兼容 GPT-2 的纯 C 实现）改写为 Swift。目标不仅是与 C 达到同等水平，更是要利用 M 系列芯片上所有可用的工具——从 CPU 与 SIMD 指令到“秘密”AMX 协处理器以及通过 Metal 使用的 GPU——将 Swift 推向极限。
起点：性能差距
将核心的 matmul_forward 函数从 C 翻译为基础 Swift 时，最初的结果十分显著。即使在 Release 配置下并移除了运行时断言，基础 Swift 实现仍比纯 C 版本慢 15 到 20 倍。
| 模型 | 令牌/秒 | 训练迭代/秒 | 相对 llm.c 的训练性能 |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Basic Swift | 0.054 | 0.014 | 7.3% |
这相当于约 2.8 Gflop/s 的性能——在 1999 年可能已经相当惊人，但对现代 LLM 工作负载而言却不可接受。主要原因被定位为 _ArrayBuffer.beginCOWMutation()。Swift 的写时复制（COW）唯一性检查即使在数组本身唯一时也会产生巨大的开销。
缩小差距：Swift 层面的优化
为了解决 COW 瓶颈，第一步是采用 MutableSpan（在 Swift 6.2 中引入），它提供了一种几乎零开销的可靠内存访问方式。虽然这提升了训练速度，但前向传播仍然缓慢，因为 Swift 没有与 C 的 -ffast-math 标志直接对应的功能，而该标志可启用融合乘加（FMA）指令。
利用宽松数学和 SIMD
通过使用 Swift-Numerics 库及其 Relaxed.multiplyAdd 函数，实现终于能够利用 fmla（SIMD 向量化 FMA）指令。仅此改动就使令牌每秒的速度提升了近 10 倍。
循环展开与 InlineArray
为了匹配经过优化的 C 实现——该实现通过跨步循环来鼓励编译器展开——作者使用了 InlineArray（Swift 6.2）。这使得可以使用栈分配的缓冲区，避免在循环中堆分配数组的高开销。在此阶段，“Fast Swift”达到了与 C 相当的水平，甚至在训练迭代每秒上略有超越（相当于 llm.c 的 106.6%）。
扩展规模：多线程与 AMX
虽然单线程性能已经解决，但下一步的提升需要利用所有可用的 CPU 核心。使用 DispatchQueue.concurrentPerform 可以将工作负载分配到 M3 Max 的 16 核上。然而，这在代码中引入了大量“视觉杂乱”，需要使用 withUnsafeMutableBufferPointer 和 @unchecked Sendable 包装器来绕过 Swift 的并发安全检查。
“秘密”武器：AMX
除了标准 SIMD，Apple Silicon 还内置了 AMX（Apple 矩阵协处理器）。虽然 Apple 仅通过 Accelerate 框架正式公开该功能，但诸如 AMX_MATFP 的逆向工程指令可以直接操作 16×16 的瓦片。
警告： 不建议在生产环境中直接使用 AMX 指令，因为它们未文档化且可能导致二进制兼容性中断。仍推荐使用 Accelerate 框架。
实现 AMX 指令后，训练性能提升至原始 llm.c 实现的 958.8%。
最终前沿：Metal 与 GPU
为了进入 Tflop/s 级别的领域，工作负载被迁移到使用 Metal 的 GPU 上。此过程包括在 Metal/C++ 中编写计算内核以及在 Swift 中编写调用层。
- Basic Metal（基础 Metal）： 一个朴素的内核相较于 AMX 提供了适度的提升。
- Threaded Metal（线程化 Metal）： 优化
threadsPerThreadgroup带来了显著的提升（相当于llm.c的 2204.6%）。 - Tiled Metal（瓦片化 Metal）： 通过实现瓦片内核以提升内存局部性（减少遍历长行的需求），性能最终突破了 1 Tflop/s 的瓶颈。
最终性能对比
| 模型 | 令牌/秒 | 训练迭代/秒 | 相对 llm.c 的训练性能 |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Multithreaded Swift | 4.356 | 1.014 | 558.5% |
| AMX | 5.884 | 1.678 | 958.8% |
| Tiled Metal | 11.123 | 5.351 | 3057.7% |
技术洞察与反思
FMA 争论
社区讨论中提出的关键点是关于使用 -ffast-math。虽然作者使用它来启用 FMA，但一些专家认为 -ffast-math 范围过宽，可能导致不期望的数值误差。对于仅想获得 FMA 而不承担完整 fast-math 风险的用户，推荐的替代方案是 -ffp-contract=fast。
GPU 软件壁垒
从“Basic Metal”到“Tiled Metal”的转变难度凸显了为何 NVIDIA 的 CUDA 等软件生态系统仍然占据主导地位。峰值 GPU 性能不仅取决于硬件本身，更依赖于针对特定数据形状的庞大且高度调优的内核库。
结论
从最初 2.8 Gflop/s 的朴素实现出发，作者最终达到了 1.1 Tflop/s——性能提升了 382 倍。此过程表明，虽然 Swift 能够匹配甚至超越 C 的速度，但往往需要以牺牲语言本身的优雅为代价，因为代码会降级为不安全指针和手动内存管理。对于生产环境的应用，教训十分明确：应使用已经经过多年优化的成熟框架（Accelerate、CoreML、MPSGraph）。