Rust in Everything
基于 burn 0.22.1

计时

主线位置：张量主线的「计算 → 读回」之间：给一次计算计时，时钟读在哪里才算量到了计算本身。
版本：PyTorch 2.15.0（CPU、MPS）、burn 0.22.1（Flex 设备、Metal 设备）
sequenceDiagram
  participant H as Python 或 Rust
  participant D as 设备
  H->>H: t0 = now
  H->>D: a @ b 或 a.matmul(b)
  H->>H: 不同步就读 t1
  Note over H: t1 - t0 只是调用开销
  H->>D: synchronize 或 sync
  D-->>H: 做完了
  H->>H: t2 = now
  Note over H: t2 - t0 才含计算

上一页说明了 GPU 上的调用为什么先返回。这一页给出两边计时的对应写法，量同一个 [1024, 1024] 的 f32 matmul，每种写法都写对的和错的两种。

要做的事PyTorchburn
等设备做完torch.mps.synchronize()（CUDA 上 torch.cuda.synchronize()）device.sync()
读回一个数，顺带同步t.item()、t.cpu()t.into_scalar()、t.into_data()
设备上的时间torch.mps.Event(enable_timing=True) 加 elapsed_timedevice.profile_with(ProfileOptions::default().flush(), ..)
自动同步的计时工具torch.utils.benchmark.Timer没有，自己写

实验代码：PyTorch 侧 exec_gpu.py 的 mm 模式；burn 侧三个 crate，exec-gpu（默认 feature 加 flex、metal）、exec-cpu（本卷模板的 feature，default-features = false）、exec-amx（默认 feature 加 flex、apple-amx），代码相同。全部 release 构建，每种写法 20 次，报最小值和中位数。平台 M2 Max、macOS 27、rustc 1.99.0；GPU 那组测量时 uptime 的 1 分钟负载约 24，CPU 那组约 7，apple-amx 那组约 19。

GPU：不同步量到的是什么

for _ in range(reps):
  t = time.perf_counter()
  c = a @ b
  nosync.append(time.perf_counter() - t)
  torch.mps.synchronize()
  synced.append(time.perf_counter() - t)
  del c

for _ in 0..reps {
  let t = Instant::now();
  let c = a.clone().matmul(b.clone());
  nosync.push(t.elapsed());
  device.sync().unwrap();
  synced.push(t.elapsed());
  drop(c);
}

设备时间的写法：

def event_time(f) -> float:
  start = torch.mps.Event(enable_timing=True)
  end = torch.mps.Event(enable_timing=True)
  start.record()
  out = f()
  end.record()
  end.synchronize()
  del out
  return start.elapsed_time(end) / 1e3

let (out, d) = device
  .profile_with(ProfileOptions::default().flush(), || a.clone().matmul(b.clone()))
  .unwrap();
ts.push(read_sync(d.resolve()).map(|t| t.duration()).unwrap_or_default());
drop(out);

同一个进程里第二次起的结果（第一次单独讲）：

...
mps 1024x1024 matmul no sync, 20 runs: min 28.166µs median 35.458µs
mps 1024x1024 matmul with sync, 20 runs: min 535.250µs median 1.118ms
mps 1024x1024 matmul mps.Event time, 20 runs: min 419.333µs median 445.333µs
mps 1024x1024 matmul + sum().item(), 20 runs: min 592.459µs median 985.416µs

...
metal 1024x1024 matmul no sync, 20 runs: min 791.000ns median 1.375µs
metal 1024x1024 matmul with sync, 20 runs: min 651.666µs median 688.791µs
metal 1024x1024 matmul device time (profile_with flush), 20 runs: min 427.375µs median 429.334µs
metal 1024x1024 matmul device time (profile, no flush): 20 of 20 read as no time
metal 1024x1024 matmul + sum().into_scalar(), 20 runs: min 706.750µs median 751.958µs
  • 不同步的那一行两边都是错的，量到的是调用开销：PyTorch 约 30 微秒（Python 一侧的调用和把 op 交给 MPS 的开销），burn 约 1 微秒（登记一条 IR）。这两个数相差 30 倍，和 matmul 本身无关。
  • 同步之后的墙钟时间两边在同一水平（0.54 对 0.65 毫秒的最小值），中位数受机器负载影响跳得多，PyTorch 的中位数到了 1.1 毫秒。
  • 设备时间两边几乎相同：0.42 对 0.43 毫秒。两边的 matmul 用的不是同一个内核（PyTorch 走 MPS 的实现，burn 走 cubecl 经 autotune 选出的内核），这个形状上速度一样。墙钟比设备时间多出的 0.1 到 0.2 毫秒是同步路径上的开销。
  • profile 不加 flush 在 Fusion 下 20 次全是 no time：matmul 还在队列里，窗口就关了。ProfileOptions::flush 的文档说它「measures the closure's whole work at the cost of running it as production would not」：强制在窗口结束时执行，代价是打断了正常运行时的攒批。
  • 读回一个数（.item() / into_scalar()）也是同步，还多了一次 sum 和拷回，比单纯同步多 0.05 毫秒左右。

torch.utils.benchmark.Timer 会替你同步，它的文档写着：「If PyTorch was built without accelerators or there is no accelerator present, this defaults to timeit.default_timer; otherwise it will synchronize accelerators before measuring the time.」。burn 0.22.1 没有对应的工具，计时代码里每个读时钟的点之前要自己 sync。另外，burn 的 Fusion 会省掉没人读的结果，被计时的结果要留到 sync 之后再丢（上面的循环里 drop(c) 在 sync 之后），否则什么都没算（基准测试测出来几乎为零）。

GPU：第一次

每个进程里第一次 matmul 单独记：调用返回的时间、同步之后的时间。连跑两个进程：

### burn metal process 1
metal: device init 16.118ms
metal 1024x1024 matmul first call: returned 10.875µs, after sync 3.600s
...
### burn metal process 2
metal: device init 16.615ms
metal 1024x1024 matmul first call: returned 12.584µs, after sync 108.038ms
...

### torch mps process 1
torch 2.15.0, threads 8
mps: device init 79.764ms
mps 1024x1024 matmul first call: returned 130.549ms, after sync 132.103ms
...
### torch mps process 2
torch 2.15.0, threads 8
mps: device init 79.016ms
mps 1024x1024 matmul first call: returned 12.708ms, after sync 14.194ms
...

两边第一次都慢，慢在不同的行：

  • PyTorch 慢在调用那一行（returned 130.549ms），同步只多了 1.5 毫秒：第一次的准备工作在调用时同步完成。第二个进程降到 12.7 毫秒，说明有一部分结果跨进程缓存了；缓存在哪一层，本页没有查。
  • burn 的调用照样 11 微秒返回，第一次的代价全落在 sync 上：冷缓存 3.6 秒（autotune 实测一批候选内核），第二个进程 108 毫秒（调优结果从磁盘读出，内核仍要编译）。这个缓存默认放在「当前目录往上最外层 Cargo.toml」旁边的 target/environment，换一个目录启动又是冷的，见 burn 卷的 GPU 上第一次很慢、计时忽快忽慢。

所以从 PyTorch 搬过来的「先跑一次热身」照样要做，但在 burn 里热身要以 sync 或读回结束，否则热身那次只是登记，编译和 autotune 会落进第一次正式计时里。

❗IMPORTANT
动脑筋：把上面 burn 的热身写成 let _ = a.clone().matmul(b.clone());（不 sync、不读回），然后开始计时。第一次正式计时的 after sync 会是多少：0.65 毫秒左右，还是秒级？冷缓存，平台和构建同上。答案在本页末尾。

CPU：两边都同步，快慢看库

CPU 上不需要同步，「不同步」和「同步之后」两个数一样。PyTorch 固定线程数（torch.set_num_threads）分别跑 1 和 8；burn 的 Flex 分三种构建跑，打开默认 feature 的两种用 RAYON_NUM_THREADS 固定线程数：

### torch cpu threads 1
torch 2.15.0, threads 1
...
cpu 1024x1024 matmul no sync, 20 runs: min 831.334µs median 863.167µs
cpu 1024x1024 matmul with sync, 20 runs: min 832.209µs median 863.750µs
...
### torch cpu threads 8
torch 2.15.0, threads 8
...
cpu 1024x1024 matmul no sync, 20 runs: min 816.209µs median 850.542µs
cpu 1024x1024 matmul with sync, 20 runs: min 816.875µs median 850.625µs
...

### burn flex template
...
flex(template) 1024x1024 matmul no sync, 20 runs: min 21.637ms median 22.352ms
flex(template) 1024x1024 matmul with sync, 20 runs: min 21.637ms median 22.353ms
### burn flex default features, RAYON_NUM_THREADS=8
...
flex 1024x1024 matmul no sync, 20 runs: min 4.113ms median 4.612ms
flex 1024x1024 matmul with sync, 20 runs: min 4.114ms median 4.613ms
...
### burn flex default features, RAYON_NUM_THREADS=1
...
flex 1024x1024 matmul no sync, 20 runs: min 21.455ms median 21.973ms
flex 1024x1024 matmul with sync, 20 runs: min 21.455ms median 21.973ms
...

exec-amx：

### RAYON_NUM_THREADS=8
flex+apple-amx 1024x1024 matmul first call: 3.004ms
flex+apple-amx 1024x1024 matmul with sync, 20 runs: min 1.430ms median 1.640ms
### RAYON_NUM_THREADS=1
flex+apple-amx 1024x1024 matmul first call: 5.127ms
flex+apple-amx 1024x1024 matmul with sync, 20 runs: min 3.477ms median 3.540ms
构建1 线程8 线程
PyTorch CPU0.83 ms0.82 ms
burn Flex，本卷模板（无 simd、rayon）21.6 ms只有单线程
burn Flex，默认 feature21.5 ms4.1 ms
burn Flex，默认 feature 加 apple-amx3.5 ms1.4 ms
  • PyTorch 的线程数从 8 降到 1，时间不变。这个 wheel 的 torch.__config__.show() 写着 BLAS_INFO=accelerate：matmul 交给苹果的 Accelerate 库，set_num_threads 管不到它。所以比 CPU 速度时，「固定了 PyTorch 的线程数」在 matmul 上不等于两边用了同样的算力。
  • Flex 的 matmul 用 gemm crate（Flex：CPU 后端）。本卷模板关了默认 feature，没有 rayon，只有一个线程；默认 feature 下 8 个线程快 5 倍多。
  • burn 有一个 apple-amx feature，在 burn-flex 里打开的是 gemm/experimental-apple-amx（名字里就写着 experimental）。打开后 8 线程 1.4 毫秒，离 PyTorch 还差 1.7 倍。burn 还有一个 accelerate feature，接到 ndarray 后端的 blas-accelerate，离线构建失败（缓存里没有 blas-src），本页没有测。

几种构建算出的结果是否一样：用确定的输入（a = sin(k)、b = cos(k)，[256, 256]）和 float64 参考比：

torch cpu            max|x - f64 ref| = 2.008e-06   max|x - torch cpu| = 0.000e+00
torch mps            max|x - f64 ref| = 2.008e-06   max|x - torch cpu| = 0.000e+00
burn flex template   max|x - f64 ref| = 2.227e-06   max|x - torch cpu| = 2.027e-06
burn flex default    max|x - f64 ref| = 2.227e-06   max|x - torch cpu| = 2.027e-06
burn flex apple-amx  max|x - f64 ref| = 2.227e-06   max|x - torch cpu| = 2.027e-06
burn metal           max|x - f64 ref| = 2.227e-06   max|x - torch cpu| = 2.027e-06
max |ref| = 1.095

burn 的四种构建彼此逐位相同，和 PyTorch 差 2.0e-6；每个元素是 256 项乘积之和、量级约 1，这个差是求和顺序不同造成的舍入差（「对得上」的标准见怎么读这一卷）。换构建、换线程数只改变速度，没有改变这一题的结果。

动脑筋答案

秒级。exec-gpu 的 warm 模式：热身一次不 sync，然后计时 3 次，每次 sync。热身分两种：结果当场丢掉（let _ = ...），或者留在变量里但不读。每次运行前删掉 crate 目录下的 target/environment（冷缓存），最后再用热缓存跑一次（负载约 3）：

let kept = if kind == "kept" {
  Some(a.clone().matmul(b.clone())) // kept alive, never synced or read
} else {
  let _ = a.clone().matmul(b.clone()); // dropped right away
  None
};
for i in 0..3 {
  let t = Instant::now();
  let c = a.clone().matmul(b.clone());
  device.sync().unwrap();
  println!("warm-up {kind}, timed run {i}: after sync {:.3?}", t.elapsed());
  drop(c);
}

### cold cache (target/environment deleted), warm-up dropped
warm-up dropped, timed run 0: after sync 3.843s
warm-up dropped, timed run 1: after sync 694.125µs
warm-up dropped, timed run 2: after sync 658.250µs
### cold cache (target/environment deleted), warm-up kept
warm-up kept, timed run 0: after sync 3.379s
warm-up kept, timed run 1: after sync 748.875µs
warm-up kept, timed run 2: after sync 661.167µs
### warm cache, warm-up dropped
warm-up dropped, timed run 0: after sync 105.686ms
warm-up dropped, timed run 1: after sync 2.016ms
warm-up dropped, timed run 2: after sync 1.904ms
  • 丢掉结果的热身什么都没做：Fusion 把没人读的结果整块省掉，autotune 的 3.8 秒全落进第一次正式计时。
  • 留住结果的热身也没用：它只是排在队列里，第一次正式计时的 sync 把它和正式的那次一起执行，3.4 秒照样算进来。
  • 热缓存下第一次仍要 106 毫秒（编译内核）。同一段代码在 PyTorch 里，热身那一行自己就会卡住、把准备工作做完（上面第一个进程 130 毫秒），不需要同步。

小结

  • GPU 上不同步读时钟，两边量到的都是调用开销：PyTorch MPS 约 30 µs，burn Metal 约 1 µs。
  • torch.mps.synchronize() 对 device.sync()，.item() 对 into_scalar()；设备时间 mps.Event 对 profile_with(.. .flush(), ..)，[1024, 1024] matmul 两边都是 0.42–0.43 ms。
  • 第一次：PyTorch 慢在调用那一行，burn 慢在 sync（冷缓存 3.6 s，autotune）；burn 的热身要以 sync 或读回结束。
  • burn 没有 torch.utils.benchmark.Timer 那样自动同步的工具；被计时的结果要留到 sync 之后再丢。
  • CPU 上两边都同步。PyTorch 的 matmul 走 Accelerate，不受 set_num_threads 影响；Flex 模板单线程 21.6 ms，默认 feature 8 线程 4.1 ms，加 apple-amx 1.4 ms；结果逐位不变。

💬 关联讨论

发起讨论
加载中...