Rust in Everything
基于 rayon 1.12.0

几十个元素的 par_iter 比串行慢几万倍

主线位置:主线 1 和主线 2 的固定成本。进池的三条路径见 in_worker:从哪进池,切分见 bridge 与自适应切分,唤醒见 唤醒策略与实测。

flowchart LR
  call["main 里 v.par_iter().sum(),16 个元素"] --> split{"len ≥ 2?"}
  split -->|"否:n = 1"| local["在调用线程上顺序算,不进池"]
  split -->|"是"| cold["根 join 走 in_worker_cold:注入"]
  cold --> wake["叫醒 worker;worker 再叫醒至多 2 个"]
  wake --> tree["池外多切一层:16 个元素切成 16 片、15 次 join"]
  tree --> set["LockLatch::set 唤醒调用线程"]
  set --> cost["12 线程约 21 µs;串行约 1 ns"]

现象

把一个很短的循环(几十到几百个元素、每个元素几纳秒的活)改成 par_iter(),不但没变快,反而慢了几个数量级;采样里调用线程大部分时间停在 in_worker_cold / __psynch_cvwait(Linux 上是 futex)。RAYON_NUM_THREADS=1 时变快不少,但仍比串行慢。

复现

实验 E13(exp-rayon/outside-call/,每项 2 万次取平均,3 次的中位数,1 分钟负载 2.78–2.98,详见 in_worker:从哪进池):

项目12 线程1 线程
从 main 调一次空 rayon::join17.2 µs3.85 µs
16 个元素 par_iter().sum(),从 main20.7 µs4.28 µs
同上,在 pool.install 里1.61 µs15 ns
16 个元素 iter().sum()1 ns1 ns

两次调用之间间隔够久、池里的线程已经睡在 Condvar 上时更贵:空闲 10 ms 后从 main 调一次空 join,12 线程 39–42 µs(唤醒策略与实测)。

池内也有门槛。E1(join 与 join_context):fib(36) 递归到底(cutoff 1,2416 万次 join)在 12 线程上只有 1.8 倍,cutoff 15–20 时约 10 倍;1 线程下每次 join 约 6.6 ns。E4(bridge 与自适应切分):10⁶ 个每元素一次乘加的轻循环,默认 0.32 ms,with_min_len(10⁴) 0.11 ms,with_max_len(10) 0.49 ms。

从哪来

一次从池外发起的并行调用,固定要付这几笔:

  1. 进池:调用线程不是 worker,根 join_context 走 in_worker_cold,注入全局队列后睡在线程本地的 LockLatch 上。
  2. 唤醒:注入时 new_injected_jobs 叫醒睡着的线程(从 0 号起),醒来的线程在 work_found 里再叫醒至多两个(唤醒策略与实测);每一级都是一次 notify_one 加 OS 调度。
  3. 切分:池外调用时根的两个孩子 migrated() 都为真,多切一层,12 线程最少 64 片;16 个元素就切成 16 片、15 次 join(bridge 与自适应切分)。只有 n = 1 时根本不切,整个计算在调用线程上完成,不进池。
  4. 回来:根任务结束时 LockLatch::set 唤醒调用线程,又一次 Condvar 往返。

16 个元素的实际计算只有一纳秒量级,这几笔加起来是微秒到几十微秒。

两个前几卷的实例都是这个问题:

  • candle 的 softmax_last_dim 对 9 行 × 33 个数用 par_chunks,「走 rayon 要几十微秒」,RAYON_NUM_THREADS=1 降到约 9 µs;pytorch-to-candle 卷的采样 217 个样本里 205 个在 in_worker_cold(candle 卷 quantized-llama、pytorch-to-candle 卷 speed;拆解见 唤醒策略与实测)。
  • burn 的 Flex 后端只在大张量上开 rayon:求和约 16 MiB、矩阵乘 M·N·K ≥ 192³ 才并行(burn 卷 flex)。这就是用阈值绕开上面的固定成本。

怎么改

  • 按数据量切换:if v.len() < THRESHOLD { v.iter().sum() } else { v.par_iter().sum() }。阈值按「整段计算明显大于几十微秒」定,从池外调用时尤其如此;本机 12 线程,从池外进池一次就是 17–40 µs。
  • with_min_len:每元素很便宜时设一个下限,让每片的工作量明显大于一次 join(E4:10⁴ 最好);设成 len 就完全不进池。别对重计算设太大,片数少于线程数会掉速(E4:min_len = 10⁵ 只有 8 片,降到 7.5 倍)。
  • 把小并行段合进一次大的 install:反复从池外调小的 par_iter,每次都付进池和唤醒;把整个循环放进 pool.install(|| ...) 里,内部的 par_iter 走池内路径,16 个元素从 21 µs 降到 1.6 µs(E13)。candle 的 Device::with_context 就是这个思路。
  • RAYON_NUM_THREADS=1 只当诊断:它确认「慢在调度、不在计算」,但同时关掉了所有并行;E13 里 1 线程从池外调用仍要 4 µs。

小结

  • 从池外进池一次是微秒级(12 线程 17–21 µs,池睡着时约 40 µs),池内 join 是纳秒级;16 个元素串行只要 1 ns(E13)。
  • 来源:in_worker_cold 的注入与 LockLatch、链式唤醒、池外多切一层、回程的 Condvar 唤醒。
  • 改法:按数据量切换串行、with_min_len、把多次小并行合进一次 install;RAYON_NUM_THREADS=1 只用来诊断。

💬 关联讨论

发起讨论
加载中...