Rust in Everything
基于 burn 0.22.1

基准测试测出来几乎为零

主线位置:算子主线的 Fusion 一层。op 登记进流后,fusion 服务端按块融合、执行;一个没人再用的输出在执行前就被丢掉,几块不相干的计算也可以合成一块。机制见 Fusion 装饰器和融合搜索。

flowchart TB
  a["A 每次丢掉结果"] --> elide["整块被省掉 0 个内核"]
  c["C 结果全留着 最后 sync"] --> merge["20 次合成 1 个内核"]
  b["B 留住结果 每次 sync"] --> one["每次 1 个内核"]
  p["profile_with flush"] --> dev["设备时间"]

实验 crate pc-gpu bench,Metal,a、b、c 都是 [4096, 4096] 的 f32,算 (a * b + c).exp()(fusion 会把它融成一个逐元素内核),每种写法 20 次,release 构建,M2 Max、macOS 27、rustc 1.99.0。用 tracing 数 cubecl 的 launch_inner span 得到启动次数(需要 burn 的 tracing feature)。正式计时前已经算过一次,编译不在里面。负载(uptime 1 分钟)约 41。

现象

// A: result thrown away, one sync at the end.
let t = Instant::now();
for _ in 0..reps {
  let _ = expr();
}
device.sync().unwrap();

// B: result kept until sync, sync every evaluation.
for _ in 0..reps {
  let t = Instant::now();
  let d = expr();
  device.sync().unwrap();
  times.push(t.elapsed());
  drop(d);
}

// C: all results kept, one sync at the end.
let t = Instant::now();
let outs: Vec<Tensor<2>> = (0..reps).map(|_| expr()).collect();
device.sync().unwrap();
4096x4096 f32, (a * b + c).exp(), 20 evaluations
A drop each result:       346.672µs per eval, 0 launch(es)
B keep + sync each:         1.378ms median,   20 launch(es)
C keep all, one sync:       4.122ms per eval, 1 launch(es)
4096x4096 f32, (a * b + c).exp(), 20 evaluations
A drop each result:       347.037µs per eval, 0 launch(es)
B keep + sync each:         1.306ms median,   20 launch(es)
C keep all, one sync:       4.065ms per eval, 1 launch(es)
  • A 有 sync,看起来是个合格的基准,可是一个内核都没启动,每次 0.35 ms 里没有任何 GPU 计算。
  • C 的 20 次计算只启动了 1 个内核,「每次 4.1 ms」并不是单次计算的时间。
  • 只有 B 的每次对应一次真实的内核执行。

出处

A:let _ = expr(); 让结果当场被 drop。fusion 登记了 Mul、Add、Exp,紧跟一个 Drop;服务端处理这一块时,输出已经没人要,整块被省掉(Fusion 装饰器「没人读的结果」一节用 2×2 的张量看过:4 个 op 的块,sync 之后启动 0 次)。burn 自己的代码里有一段注释说的就是这件事,在 capture 的预热循环里(graph.rs):

        // Hold the output alive across the sync. A lazy backend (fusion) elides
        // a computation whose result is dropped before it drains — so dropping
        // the warmup output would skip the very kernels (and allocations) the
        // capture must reuse, and the capture run would then be the first to
        // allocate them, faulting mid-capture. Keeping `out` until after `sync`
        // forces the closure to actually execute and populate the pool.
        let out = closure();
        let _ = Dispatch::sync(dispatch);
        drop(out);

C:20 次 expr() 在一次 sync 前登记,融合搜索会把互不相干的块继续合并,只要融合器还开着(融合搜索「不相干的计算也会被合并」)。这里实测只启动了 1 个内核,20 个结果都由它写出(融合搜索那一页用日志看过同类情况:5 条链合成一个读 5 个输入、写 5 个输出的内核)。它的耗时既不是单次的 20 倍,也不能拿来除以 20 当单次。

Flex 是即时执行的,没有这两个问题;同一段基准代码换到 Metal 上,量的东西就变了。

修法

按 B 写:每次计算的结果留到 sync 之后再丢,每次单独 sync。这样量到的是墙钟时间,包含提交和等待的开销。

量设备时间:Device::profile_with 配 ProfileOptions::default().flush()。flush 不能少,Device::profile 的文档写着 fusion 会把闭包最后的 op 留在队列里,「a window over lazy work alone can read as empty」:

let (out, d) = device.profile(|| (a.clone() * b.clone() + c.clone()).exp()).unwrap();
show(&format!("[{i}] profile expr"), d);
drop(out);
let (out, d) = device
  .profile_with(ProfileOptions::default().flush(), || (a.clone() * b.clone() + c.clone()).exp())
  .unwrap();
show(&format!("[{i}] profile_with(flush) expr"), d);
drop(out);
[0] profile expr                   no time (device)
[0] profile_with(flush) expr        786.500µs (device)
[0] profile_with(flush) matmul       23.654ms (device)
[1] profile expr                   no time (device)
[1] profile_with(flush) expr        761.417µs (device)
[1] profile_with(flush) matmul       22.694ms (device)
[2] profile expr                   no time (device)
[2] profile_with(flush) expr        758.666µs (device)
[2] profile_with(flush) matmul       22.549ms (device)

(show 用 burn::tensor::read_sync(d.resolve()) 取出 ProfileDuration,None 打成 no time。)同一个表达式,设备时间约 0.76 ms,B 的墙钟中位数约 1.3 ms,多出来的部分花在内核之外。

如果确实想量「一批互不相干的计算一起跑」的吞吐量,C 的写法是对的,但要按批报告,并说明它被合成了几个内核。

怎么确认没中招

  • 基准里装一个 tracing 层数 launch_inner span(cubecl 每次启动内核一个),看启动次数是不是你以为的次数。0 次说明结果被丢了,远少于计算次数说明被合并了。
  • 被计时的结果在 sync 之后才 drop;let _ = ... 和不绑定变量的表达式语句都是当场 drop。
  • 用 profile 量 fusion 下的计算时,加 ProfileOptions::default().flush();结果是 None 就说明窗口里什么都没执行。

小结

  • fusion 下,没人读、执行前就被丢掉的结果,相应的计算不会执行;基准量到的只是登记。
  • 一次 sync 前的多次独立计算会被合成一个多输出内核,总时间不能拆成单次。
  • 每次留住结果并 sync,或用 Device::profile_with(ProfileOptions::default().flush(), ..) 量设备时间;用启动次数核对。

💬 关联讨论

发起讨论
加载中...