Rust in Everything
基于 tenferro-runtime 0.8.0

每次调用都新建 engine

主线位置:直接计算主线的第 2 步,let mut engine = CpuEngine::...。engine 是调用方持有的值,带着线程 team、缓冲池和计划缓存;它的构造代价逐项测量见 CpuEngine。

flowchart TB
  f["fn step 每次调用"] --> b["Pool::builder build"]
  b --> c{"默认设置吗"}
  c -->|是| hit["从进程缓存拿线程"]
  c -->|否| sp["新建 8 个线程"]
  hit --> res["预留 team"]
  sp --> res
  res --> e["空缓冲池 空计划缓存"]
  e --> mm["matmul"]
  mm --> d["drop 引擎"]
  fix["engine 建一次"] --> pass["可变借用传下去"]

实验 crate exp-tenferro8/pit1-pit,bin p07_engine_per_call,release;f64 方阵 a.matmul(&a, &mut engine),先热身 20 次,32×32 计 500 次、256×256 计 100 次,取最小值和中位数;连跑三轮,下面引第三轮,三轮的最小值相差在 30% 以内。uptime 1 分钟负载 3 到 4(有别的写手在跑)。M2 Max、macOS 27、rustc 1.99.0。

现象

0.8.0 去掉了全局的后端会话,运算都要一个 &mut CpuEngine。最顺手的写法是在用到的函数里就地建一个:

fn step(a: &TypedTensor<f64>) -> anyhow::Result<TypedTensor<f64>> {
  let mut engine = CpuEngine::new(Pool::builder().threads(8).build()?);
  Ok(a.matmul(a, &mut engine)?)
}

能算对,只是慢。和复用同一个 engine 对比:

== 32x32 f64 matmul, 500 reps after 20 warm-up
reused CpuEngine::serial()                   min      2.71 us  median      2.83 us
new CpuEngine::serial() per call             min      7.04 us  median      7.58 us
reused 8-thread engine                       min      3.42 us  median      3.54 us
new 8-thread engine per call                 min     17.21 us  median     23.12 us
new 8-thread engine, stack_size(8 MiB), per call min    168.29 us  median    227.42 us
== 256x256 f64 matmul, 100 reps after 20 warm-up
reused CpuEngine::serial()                   min    683.04 us  median    720.54 us
new CpuEngine::serial() per call             min    704.75 us  median    732.67 us
reused 8-thread engine                       min    117.33 us  median    123.25 us
new 8-thread engine per call                 min    140.75 us  median    200.88 us
new 8-thread engine, stack_size(8 MiB), per call min    297.58 us  median    369.62 us

小矩阵上每次新建 8 线程 engine,一次 matmul 慢 5 倍;给池设了 stack_size(或者 thread_name、on_start、qos),每次多出约 0.17 ms,256×256 时比单线程复用还只快一倍。程序里这样的调用成千上万次时,差距就是线程池的建和拆。

还有一个看不见的代价:输出缓冲区回不到下一次调用。同一个程序的结尾:

reused engine pool after one dropped output: BufferPoolStats { buffers: 1, capacity_bytes: 32768 }
fresh engine pool:                           BufferPoolStats { buffers: 0, capacity_bytes: 0 }

复用的 engine 丢掉一个 64×64 的结果后,32 KiB 的缓冲区回到了它的池里,下一次同样大小的输出直接拿走;新 engine 的池永远是空的。

出处

CpuEngine 的文档注释说它「owns one threadteam::Pool (or none for CpuEngine::serial) and the workspace operations write: the buffer pool that recycles outputs, the contraction scratch, and a handle to a CpuPlanCache」。新建一个 engine,缓冲池和计划缓存都从空的开始,这是串行 engine 每次多出约 4 µs 的地方:CpuEngine::serial() 本身不到 1 µs(见 CpuEngine)。

多线程 engine 还要建池、预留 team。threadteam 0.2 会缓存默认设置的池的线程,build_checked 里先判断能不能复用:

        let shared = self.thread_name.is_none()
            && self.on_start.is_none()
            && self.stack_size == DEFAULT_STACK_BYTES
            && self.qos.is_none();
        let workers = match shared.then(|| cache::take(n)).flatten() {

默认设置的池在 drop 时把线程放回按宽度索引的进程缓存,下一个同宽度的池直接拿走,所以「每次新建 8 线程 engine」剩下的主要是预留 team 的十几微秒。设了 stack_size 等任何一项,shared 为假,每次都 spawn 新线程、drop 时 join,于是有 0.17 ms 的差。默认栈已经是 16 MiB,通常没有理由再设。

0.7.1 用的是 CpuBackend 加会话的模型,这个坑当时叫 backend-per-call。会话换成调用方持有的 engine 的经过见 CpuEngine。

修法

engine 在程序入口建一次,以 &mut 往下传;需要并发就每个线程各建一个(CpuEngine 是 Send):

fn step(a: &TypedTensor<f64>, engine: &mut CpuEngine) -> anyhow::Result<TypedTensor<f64>> {
  Ok(a.matmul(a, engine)?)
}

let mut engine = CpuEngine::new(Pool::builder().threads(8).build()?);
for _ in 0..steps {
  let c = step(&a, &mut engine)?;
  // ...
}

上面的「reused 8-thread engine」一行就是这种写法:32×32 一次 3.4 µs,256×256 一次 117 µs。小矩阵上串行 engine 反而更快(2.7 µs),线程数怎么选见 12 个线程比 8 个慢。

怎么确认没中招

  • 搜代码里的 CpuEngine::new(、CpuEngine::serial( 和 Pool::builder():它们应该只出现在程序入口、线程入口或测试的开头,不在循环体和每次调用的函数里。
  • 在热路径上打印 engine.buffer_pool_stats():跑过几步之后 buffers 还是 0,说明每步用的都是新 engine。
  • Pool::builder() 上的 stack_size、thread_name、on_start、qos 会让池不进缓存;只在建一次的 engine 上用它们。

小结

  • 新 engine 的缓冲池和计划缓存都是空的;多线程 engine 每次还要建池、预留 team。32×32 的 matmul 每次新建 8 线程 engine 从 3.4 µs 变成 17 µs。
  • threadteam 只缓存默认设置的池;设了 stack_size 之类的池每次新建线程,多出约 0.17 ms。
  • engine 建一次,以 &mut 传下去;并发时每个线程一个。

💬 关联讨论

发起讨论
加载中...