基准测试测出来几乎为零
主线位置:算子主线的
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_innerspan(cubecl 每次启动内核一个),看启动次数是不是你以为的次数。0 次说明结果被丢了,远少于计算次数说明被合并了。 - 被计时的结果在
sync之后才 drop;let _ = ...和不绑定变量的表达式语句都是当场 drop。 - 用
profile量 fusion 下的计算时,加ProfileOptions::default().flush();结果是None就说明窗口里什么都没执行。
小结
- fusion 下,没人读、执行前就被丢掉的结果,相应的计算不会执行;基准量到的只是登记。
- 一次
sync前的多次独立计算会被合成一个多输出内核,总时间不能拆成单次。 - 每次留住结果并
sync,或用Device::profile_with(ProfileOptions::default().flush(), ..)量设备时间;用启动次数核对。