每次调用都新建 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传下去;并发时每个线程一个。