Rust in Everything
基于 burn 0.22.1

GPU 上第一次很慢、计时忽快忽慢

主线位置：算子主线的后半段。a.matmul(b) 在 Metal 上只登记一条 IR 就返回，内核在 fusion 服务线程上、读数据或 sync 时才提交；第一次遇到一个形状还要编译内核、跑 autotune。这几步分别见一次 matmul、Fusion 装饰器和 autotune。

flowchart TB
  callop["matmul 返回 约 10 µs"] --> sync["sync 才等到执行"]
  sync --> first{"这个形状第一次?"}
  first -->|冷缓存| tune["编译 + autotune 几秒"]
  first -->|热缓存| comp["只编译 约 100 ms"]
  first -->|进程内第二次起| run["只执行 约 1.3 ms"]
  tune --> db["写 target/environment/default.db"]

实验 crate pc-gpu，Device::metal(DeviceKind::DefaultDevice)，[1024, 1024] 的 f32 matmul，release 构建，M2 Max、macOS 27、rustc 1.99.0。测量时机器上还有别的编译任务，uptime 的 1 分钟负载在 30 到 110 之间，每段输出前注明。

现象

一个常见的写法：读两次时钟夹住调用。

for i in 0..reps {
  let t = Instant::now();
  let c = a.clone().matmul(b.clone());
  let returned = t.elapsed();
  device.sync().unwrap();
  let synced = t.elapsed();
  drop(c);
  println!("run {i}: matmul returned after {returned:?}, after sync {synced:?}");
}

先删掉 crate 目录下的 target/environment（冷缓存），连跑两个进程（负载约 34）：

### process 1: cold (target/environment deleted)
1024x1024 f32, device init + 2 random + sync: 65.144791ms
run 0: matmul returned after 18.208µs, after sync 3.262425208s
run 1: matmul returned after 13.584µs, after sync 21.12425ms
run 2: matmul returned after 12.667µs, after sync 11.786417ms
run 3: matmul returned after 12.834µs, after sync 13.444459ms
run 4: matmul returned after 13.333µs, after sync 23.162458ms
run 5: matmul returned after 13.084µs, after sync 11.20825ms
run 6: matmul returned after 12.542µs, after sync 10.418125ms
run 7: matmul returned after 7.291µs, after sync 5.4515ms
### process 2: warm
1024x1024 f32, device init + 2 random + sync: 114.404834ms
run 0: matmul returned after 34.833µs, after sync 112.415583ms
run 1: matmul returned after 13.333µs, after sync 27.840875ms
run 2: matmul returned after 12.709µs, after sync 4.869667ms
run 3: matmul returned after 11.417µs, after sync 3.444875ms
run 4: matmul returned after 11.291µs, after sync 7.11925ms
run 5: matmul returned after 12.209µs, after sync 3.903209ms
run 6: matmul returned after 5.667µs, after sync 2.202458ms
run 7: matmul returned after 10.5µs, after sync 2.121375ms

三个现象叠在一起：

  1. matmul 总是十几微秒就返回。不 sync 读时钟，量到的是登记一条 IR 的时间。循环里连做 10 次也一样（热缓存，负载约 32）：

    10 matmuls 1024x1024: loop returned after 52.417µs, after sync 24.596042ms

  2. 第一次很慢，慢多少取决于缓存：冷缓存 3.26 秒（autotune 实测一批候选，外加上万次测启动开销的内核），热缓存 112 毫秒（调优结果从磁盘读，内核还要编译；wgpu<msl> 的编译结果只存在进程内存里，见 wgpu 与 Metal）。

  3. 之后的每次在 2 到 28 毫秒之间跳。负载更高时（约 110）连跑 20 次，很多次正好落在 10 ms、20 ms 附近：
    run 5: matmul returned after 13.167µs, after sync 19.965042ms
run 6: matmul returned after 13.875µs, after sync 9.97225ms
run 7: matmul returned after 13.292µs, after sync 20.016917ms
run 8: matmul returned after 14.041µs, after sync 19.944375ms
...
run 10: matmul returned after 17.75µs, after sync 2.060583ms

    同一个内核用设备时间量（下面的修法），20 次是 min 1.322ms median 1.331ms。墙钟里多出来的是 sync 这一路上的线程唤醒和调度，CPU 越忙越明显；本页没有追到 10 ms 的整数倍从哪一层来。

还有一个不那么显眼的：同一个二进制，换一个当前目录运行，又变回冷的。在另一个 crate 目录 pc-shuffle 下运行，然后回到 pc-gpu（负载约 45）：

### cwd = pc-shuffle (no target/environment yet)
1024x1024 f32, device init + 2 random + sync: 136.442416ms
run 0: matmul returned after 31µs, after sync 3.987680833s
run 1: matmul returned after 12.209µs, after sync 20.067584ms
run 2: matmul returned after 11.667µs, after sync 10.025334ms
default.db
default.db-tshm
default.db-wal
### cwd = pc-gpu again
1024x1024 f32, device init + 2 random + sync: 148.440916ms
run 0: matmul returned after 31.375µs, after sync 62.77325ms

出处

  • 登记即返回：Metal 上的 Device 走 dispatch<fusion<cubecl<wgpu<msl>>>>，Fusion 只把 op 记进流，见 Fusion 装饰器。Device::sync 才把流执行完并等设备做完。
  • 第一次：autotune 对每个调优键第一次出现时实测候选，结果写进 turso 库 default.db；吞吐量（启动开销、计算和带宽峰值）另存一张表，同一台机器上共用，细节见 autotune。
  • 缓存位置：默认的 CacheConfig::Target 从当前目录往上找最外层的 Cargo.toml，用它旁边的 target/environment；一个都找不到时用用户缓存目录（CacheConfig::root，cubecl 0.11.3）。它不看 CARGO_TARGET_DIR，也不看可执行文件在哪：
                // The *outermost* Cargo.toml, not the first one found walking
                // up: a workspace member has its own manifest, but `target/`
                // lives at the workspace root, so stopping at the innermost
                // manifest would scatter caches into `member/target`.
                let mut root = None;
                let mut dir = start.as_path();
                loop {
                    if let Ok(true) = std::fs::exists(dir.join("Cargo.toml")) {
                        root = Some(dir);
                    }

所以 cargo run 在 crate 目录里是热的，把二进制拷到别处运行、或者从别的目录启动，就是另一份缓存。

修法

读时钟之前 sync（或读出结果），并且把第一次单独算：

let _ = a.clone().matmul(b.clone()).into_data(); // compile + autotune out of the way
device.sync().unwrap();
let t = Instant::now();
let outs: Vec<Tensor<2>> = (0..reps).map(|_| a.clone().matmul(b.clone())).collect();
device.sync().unwrap();
let with_sync = t.elapsed();

结果要留到 sync 之后再丢，否则 fusion 可能根本不执行它，见 基准测试测出来几乎为零。

要比较内核本身，用设备时间。 Device::profile_with 配 ProfileOptions::default().flush()，在设备上打时间戳，不受 CPU 调度影响：

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);

1024x1024 f32 matmul device time, 20 samples: min 1.322ms median 1.331ms

不加 flush() 的 profile 在 fusion 下可能什么都量不到：Device::profile 的文档说 fusion 会把闭包最后的 op 留在队列里，「a window over lazy work alone can read as empty」。实测 (a * b + c).exp() 用 profile 三次都是 no time (device)，加 flush() 后是 0.76 ms 左右。

把缓存位置固定下来。 在进程开始、任何 GPU 操作之前设置 cubecl 的配置：

fn pin_cache(dir: &str) {
  use cubecl::config::{CubeClRuntimeConfig, RuntimeConfig, cache::CacheConfig};
  let mut config = CubeClRuntimeConfig::default();
  config.environment.path = CacheConfig::Directory(dir.into());
  CubeClRuntimeConfig::set(config);
}

从 pc-shuffle 目录运行一次，再从 pc-gpu 目录运行一次，第二次是热的：

### cwd pc-shuffle, PIN_CACHE=~/exp-burn/pc-dl-dir/pinned
256x256 f32, device init + 2 random + sync: 146.523375ms
run 0: matmul returned after 888.542µs, after sync 4.648786s
run 1: matmul returned after 10.75µs, after sync 3.080875ms
### cwd pc-gpu, PIN_CACHE=~/exp-burn/pc-dl-dir/pinned
256x256 f32, device init + 2 random + sync: 127.573ms
run 0: matmul returned after 35.334µs, after sync 67.656709ms

也可以写配置文件：cubecl.toml 里 [environment] 下 path = { directory = "/绝对路径" }（实测有效；源码里另有 "global"，用用户缓存目录）。注意配置文件同样是从当前目录往上找的，换目录运行时它也可能找不到。CubeClRuntimeConfig::set 不受这个影响，代价是它整个替代了配置：get 才会读文件、套用 CUBECL_* 环境变量，set 进去的值两样都不经过。

部署时，可以在目标机器上用同样的形状跑一遍预热，把缓存文件放在固定的位置。

怎么确认没中招

  • 计时代码里，每个读时钟的点之前有 sync 或一次读回；被计时的结果在 sync 之后才丢。
  • 第一次单独报告，并写明冷还是热：看运行目录上方最外层 Cargo.toml 旁的 target/environment/default.db 在不在。
  • 比较两个实现时用 profile_with(.. .flush(), ..) 的设备时间；墙钟时间同时报告机器负载。

小结

  • GPU op 在调用处只登记，不 sync 读到的是微秒级的登记时间。
  • 第一次 = 编译 + autotune：冷缓存几秒，热缓存约 100 ms；之后才是内核本身。
  • autotune 缓存默认在「当前目录往上最外层 Cargo.toml」旁的 target/environment，换目录就冷；用 CubeClRuntimeConfig::set 或 cubecl.toml 固定位置。
  • 墙钟受 CPU 调度影响，比内核用 Device::profile_with 加 flush() 的设备时间。

💬 关联讨论

发起讨论
加载中...