背景故事大约两年前作者翻出 21 世纪初用 C 编写的图像识别器代码想让其重新运行但失败。2024 年初围绕 LLM 讨论多但没人在 Mac 上用 Swift 训练神经网络。作者尝试 Python 库觉得无法掌控。一个月后Andrej Karpathy 发布 llm.c作者将其重写为 Swift 代码最初实现非常慢。llm.c机器学习是将模型权重应用于输入数据计算误差梯度并更新权重。矩阵乘法在机器学习中占大量工作作者关注 llm.c 中的 matmul_forward 函数一次训练迭代约有 0.2 万亿次浮点运算。纯 C 代码在 Swift 包中运行每 7 秒完成一次训练迭代推理速度不到每秒 1 个令牌。基础 Swift 实现作者尽力让基础 Swift 版本与 C 版本保持一致设置移除数组索引的运行时检查在 Release 配置下运行。但 Swift 实现极其缓慢比 C 代码慢 15 到 20 倍性能约为 2.8 Gflop/s。Span、Egg 和 Span通过 Instruments 检查发现性能开销最大的是 _ArrayBuffer.beginCOWMutation()。Swift 6.2 提供 MutableSpan 解决方案使用后训练迭代速度提高 3 倍多。轻松氛围前向传播中最耗时的是 value inp[bt * C i] * weight[o * C i]。C 有 -ffast-math 优化标志可使用融合乘加指令。Swift 没有该标志只能单独进行乘法和加法运算。使用 Swift-Numerics 的 Relaxed 实现快速数学运算后每秒处理的令牌数几乎提高 10 倍但训练性能仍比 C 慢 15%。循环展开C 语言中矩阵乘法实际函数会以 8 步为单位遍历外层循环编译器实现 8 次循环展开。Swift 6.2 提供 InlineArray 特性实现与 C 相似的循环展开C 和 Swift 推理速度基本相同Swift 训练速度略快。多线程实现llm.c 代码中有 OpenMP 注释但 Swift 包管理器使用的常规 clang 编译输出中不会生效。Swift 没有简单方法标记循环进行并行处理使用 DispatchQueue.concurrentPerform 进行并发处理。将此模式应用于训练迭代中最耗时的四个循环后速度提高 5.4 倍。绝密技巧Apple Silicon 包含 AMX 单元苹果未正式命名访问其指令的唯一公开方式是通过 Accelerate 框架。有人逆向工程了 AMX 单元工作原理作者基于此编写更快的矩阵乘法实现训练速度又提高 1.67 倍。最大功耗方案用于矩阵乘法的 Metal 代码分为内层内核用 Metal/C 编写和外层调用机制在 Swift 端。