Rust in Everything
基于 candle-core 0.11.0

Metal 上计时几乎为零

主线位置：算子主线在 Metal 上的后半段。MetalStorage::matmul 把一次 GEMM 编码进当前命令缓冲区就返回新张量，GPU 这时还没开始算；命令缓冲区每 50 次取编码器提交一次，读回或 synchronize 时才等它完成。提交和等待的细节见 Metal 后端。

sequenceDiagram
  participant U as 计时代码
  participant C as Commands
  participant G as GPU
  U->>C: matmul 编码 约 6 µs
  U->>U: 读时钟 只量到编码
  U->>C: synchronize
  C->>G: commit
  G-->>C: 约 2 ms 后完成
  C-->>U: 返回 再读时钟

实验 crate exp-candle/pa-pit，bin p10_timing，release，系统着色器缓存命中（热）；M2 Max、macOS 27、rustc 1.99.0。两个 (2048, 2048) 的 f32 随机矩阵，事先 synchronize 过。

现象

照 CPU 的习惯写一个基准：

let reps = 20;
let t = Instant::now();
let mut outs = Vec::new();
for _ in 0..reps {
  outs.push(a.matmul(&b)?);
}
let naive = ms(t);

再写一个结果用完就丢的版本，和一个读时钟前先 synchronize 的版本：

### release, CANDLE_METAL_COMPUTE_PER_BUFFER=unset
first matmul: returned after 2.673 ms, synchronized after 9.2 ms
naive:  20 matmuls 2048x2048 f32 in 6.835 ms -> 0.3418 ms each
        then reading one value back took 48.7 ms (sum -136959.8)
naive, dropped: 20 matmuls in 0.122 ms -> 0.0061 ms each
        then synchronize took 40.9 ms
synced: round 0: 20 matmuls in 49.6 ms -> 2.479 ms each, 6930 GFLOP/s
synced: round 1: 20 matmuls in 39.8 ms -> 1.992 ms each, 8626 GFLOP/s
synced: round 2: 20 matmuls in 39.3 ms -> 1.967 ms each, 8734 GFLOP/s
  • 结果用完就丢的循环，每次 matmul 6 µs，按这个数算是 2800 TFLOP/s，显然不对；之后的一次 synchronize 等了 40.9 ms。
  • 保留结果的循环每次 0.34 ms，看起来像是真的在算，但同样不是：之后读回一个值等了 48.7 ms。它和丢弃版本唯一的差别是结果还被 outs 引用着，池里没有可复用的缓冲区，每次都要新分配一个 16 MB 的输出缓冲区，多出来的就是这部分。
  • 先 synchronize 再读时钟，每次约 2 ms，8.6 到 8.7 TFLOP/s。第一轮慢一些，其余两轮稳定。

把提交频率改成每个 op 一次，结论不变：

### release, CANDLE_METAL_COMPUTE_PER_BUFFER=1
first matmul: returned after 2.741 ms, synchronized after 5.4 ms
naive:  20 matmuls 2048x2048 f32 in 8.219 ms -> 0.4110 ms each
        then reading one value back took 33.5 ms (sum -136959.8)
naive, dropped: 20 matmuls in 0.369 ms -> 0.0185 ms each
        then synchronize took 39.6 ms
synced: round 0: 20 matmuls in 42.6 ms -> 2.128 ms each, 8075 GFLOP/s
...

commit 只是把命令缓冲区交给 GPU，不等它执行完；每个 op 多了十几微秒的提交开销，GPU 时间仍然落在后面的等待里。

出处

MetalStorage::matmul 编码完 call_mlx_gemm 就返回（Metal 后端「一个 op 怎么编码」一节）。提交由 Commands::command_encoder 里的计数决定，默认数到 50 才 commit，20 次 matmul 一次都没提交，全部攒在当前命令缓冲区里。

等待只发生在两个地方：

丢掉结果不会让 GPU 少做事：candle 没有惰性求值，op 在调用时已经编码好了，40.9 ms 的等待就是 20 次 matmul 的执行时间。（burn 的 fusion 后端相反，丢掉的结果可能根本不执行，见 burn 卷的这一页。）

第一次调用另算：上面热缓存的第一次 matmul 编码就用了 2.7 ms（含从系统缓存取库、建管线），冷缓存时要几秒，见 Metal 第一次运行要好几秒。

修法

读时钟之前 synchronize，第一次单独算：

let c = a.matmul(&b)?;           // first call: library + pipeline
dev.synchronize()?;
drop(c);

let t = Instant::now();
let mut outs = Vec::new();
for _ in 0..reps {
  outs.push(a.matmul(&b)?);
}
dev.synchronize()?;
let synced = ms(t);

几条规则：

  • 每个读时钟的点之前都 synchronize()，或者读回一个值。计时区间之前也要 synchronize，否则会把之前排着的工作算进来。
  • 报告时写清楚冷热、形状、dtype、重复次数，第一次和之后分开。
  • 只量单个小 op 时，synchronize 本身的开销（提交 + 等待通知）可能比 op 还大；把多次 op 放在一个区间里再除。
  • 墙钟时间受机器负载影响；多跑几轮看是否稳定（上面第一轮就慢了 25%）。

怎么确认没中招

  • 用吞吐量做合理性检查：本页同步后实测 8.7 TFLOP/s，算出来比这个高出几个数量级（丢弃版本是 2800 TFLOP/s），一定是没等。
  • 计时循环里有 matmul 却没有 synchronize / to_*，就是本页的情况。
  • 程序里第一次读回特别慢（几十毫秒），通常不是读回慢，而是前面所有 op 的执行时间都算到了它头上。

小结

  • Metal 上的 op 只编码就返回，GPU 在之后的 synchronize 或读回时才被等待；不等就读时钟，量到的是编码时间（2048² matmul 6 µs 对实际 2 ms）。
  • 丢掉结果不会跳过计算，时间算在下一次等待头上；保留结果多出的是缓冲区分配。
  • 读时钟前 synchronize，第一次单独报告并写明冷热。

💬 关联讨论

发起讨论
加载中...