Rust in Everything
基于 tenferro-runtime 0.8.0

计时

主线位置：变换主线的「编译、准备 → 执行」两步怎么计时；TRG 主线上，20 步 TRG 的一次求值、一次求值加梯度各花多久、花在哪。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native）
flowchart TB
  jc["jit 函数调用"] --> disp["派发后立即返回"]
  disp --> bur["block_until_ready"]
  bur --> jt["真正的时间"]
  first["第一次调用"] --> trace["追踪 降级"]
  trace --> xla["XLA 编译"]
  tr["tenferro run"] --> rb["结果读回主机"]
  rb --> tt["真正的时间"]
  eng["engine 宽度"] --> tt
  load["机器负载"] --> tt

比较两边的速度，先要保证量的是同一件事。JAX 一侧最常见的错是没等结果、把第一次调用算进稳态；tenferro 一侧是没写 engine 的宽度、把 8 线程的 engine 放在已经满载的机器上。本页把这些逐个量出来，最后看 TRG 的时间花在哪里。

实验代码：JAX 一侧 timing_jax.py、ad_vec_timing.py、svd_time.py；tenferro 一侧 bin trg_staged、ad_vec_timing、svd_time（release）。M2 Max（8 个性能核加 4 个能效核）、macOS 27、rustc 1.99.0。同机有别的写作任务在跑，负载均值在 4 到 97 之间变化，每段输出都附上当时的值。TRG 都是 χ = 16、20 步、截断数 [2, 4, 4, 16, ...]，JAX 用和 tenferro 默认规则相同的 custom_jvp SVD（简并奇异值）。

JAX：调用返回不等于算完

JAX 的调用是异步派发的：fn(b) 把计算交给运行时就返回一个还没算完的数组，block_until_ready 才等它。

fn = jax.jit(jax.grad(make_trg(ks, lorentz_svd)))
t0 = time.perf_counter()
out = fn(b)
t_disp = time.perf_counter() - t0        # 返回用了多久
jax.block_until_ready(out)
t_first = time.perf_counter() - t0       # 结果真正可用

# 负载均值 9.7
value           first call: returned after   203.5 ms, ready after   389.0 ms | steady: returned after   0.060 ms, ready after   160.7 ms
grad            first call: returned after   872.6 ms, ready after  1226.6 ms | steady: returned after   0.079 ms, ready after   223.5 ms
value_and_grad  first call: returned after   819.7 ms, ready after  1166.0 ms | steady: returned after   0.085 ms, ready after   222.6 ms

稳态下 fn(b) 60 到 85 µs 就返回，真正的时间是 161 ms 和 223 ms。不 block_until_ready 地计时，量到的是派发的开销，快了三千倍。循环里如果下一次调用依赖上一次的结果，等待会隐式发生在下一次调用里，单次计时的数字就落在了错误的那一次上。

tenferro 没有这个问题需要提防：本卷实验代码计时的是 Runtime::run 加上把输出读成主机上的 Tensor（into_tensor）、再取 as_slice，返回时数据已经在手里。和 JAX 对照时，JAX 一侧就要对应地等到 block_until_ready。

第一次调用

JAX 的第一次调用里有追踪、降级和 XLA 编译。用 AOT 接口在一个新的函数对象上拆开：

f2 = jax.jit(jax.grad(make_trg(ks, lorentz_svd)))
lo = f2.lower(b)        # 追踪加降级
co = lo.compile()       # XLA 编译
co(b)                   # 第一次执行

value           AOT on a fresh function: trace+lower    44.4 ms, compile   138.4 ms, first run   184.5 ms
grad            AOT on a fresh function: trace+lower   281.1 ms, compile   525.9 ms, first run   346.7 ms
value_and_grad  AOT on a fresh function: trace+lower   287.9 ms, compile   541.3 ms, first run   350.1 ms

grad 的第一次调用 1.2 s，其中追踪加变换加降级 0.28 s、编译 0.53 s，剩下是第一次执行（347 ms，比稳态的 223 ms 慢）。

tenferro 分段的 TRG 把同样的工作拆成 5 种步各自的小程序（分段），建图、ad.grad、编译、准备一次做完：

# trg_staged，1 线程，负载均值 4.7
build+grad+compile+prepare all programs 27.4 ms (5 step shapes, 1330 ops in total)
beta=0.30 lnZ/N=0.790559063800966 exact=0.790559070951263 err=-7.150e-9
  run ms: value first 365.68 steady 336.906; value+grad first 937.81 steady 933.257

准备只要 27 ms，第一次求值比稳态慢 30 ms。tenferro 没有 XLA 那样的代码生成（记图与 jit 的「没有蠢问题」），所以「第一次」便宜，但前提是图要切小：整张 20 步图的 ad.grad 根本做不完（求导变换的代价）。

两边的稳态（同一时段，负载 5–10）：

每个 βJAX jittenferro 分段，1 线程
求值160.7 ms329–347 ms
求值加梯度222.6 ms898–933 ms
第一次（含准备）1166 ms（value_and_grad）27 ms 准备，938 ms 第一次求值加梯度

求值慢一倍的原因在最后一节。梯度慢四倍，是因为分段的反向程序把每一步的前向（包括两次 SVD）重算一遍，相当于 JAX 里给每一步套 jax.checkpoint，见 分段 和 checkpoint。

最小值还是中位数

变换主线的 jit(grad(f))，n = 1000，各跑 2000 次、重复三轮，负载均值 8.2：

# JAX
torch imported=False trial 0: min   10.2 us  median   17.2 us  p90   19.0 us
torch imported=False trial 1: min   10.0 us  median   16.9 us  p90   18.3 us
torch imported=True trial 0: min    9.6 us  median   17.0 us  p90   18.3 us
# tenferro，准备好的梯度程序
width=1 first run 5485.1 us
width=1 trial 0: min   14.6 us  median   14.9 us  p90   15.5 us
width=8 first run 443.0 us
width=8 trial 0: min   19.5 us  median   20.4 us  p90   24.7 us
  • JAX 的最小值 10 µs、中位数 17 µs，差了七成；tenferro 单线程只差 2%。按最小值比，JAX 快三成；按中位数比，tenferro 单线程快一成。所以两个数都要报，并说清楚用的是哪个。本卷其他页的表各自写明了取法（例如 eager 还是 traced 用中位数）。
  • 同一个进程里先 import torch（并设成 1 线程）对 JAX 的时间没有影响。
  • tenferro 进程里的第一次执行 5.5 ms，是一次性初始化，不能算进稳态；8 线程 engine 的第一次 0.44 ms。

线程数和负载

tenferro 的并行宽度是 engine 的属性：CpuEngine::serial() 或 CpuEngine::new(Pool::builder().threads(8).build()?)，建好就固定（12 个线程比 8 个慢）。JAX 的 CPU 后端默认自己决定。同一个 TRG：

# JAX，XLA_FLAGS='--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1'，负载均值 7.8
value           ... steady: returned after   0.060 ms, ready after   160.3 ms
value_and_grad  ... steady: returned after   0.075 ms, ready after   223.2 ms
# tenferro，8 线程，负载均值 5.8 到 8.4
beta=0.30 ...
  run ms: value first 287.78 steady 249.274; value+grad first 696.23 steady 646.370
  • 关掉 XLA 的多线程，JAX 的时间一点没变（160.3 对 160.7 ms，223.2 对 222.6 ms）。这个 TRG 的时间在 SVD 上（最后一节），单独量的 SVD 在两种设置下也一样（5.37 对 5.48 ms），所以这两个开关管不到它；LAPACK 内部是否用了多个线程，本页没有查。
  • tenferro 8 线程比 1 线程快约 30%（250 对 337 ms，646 对 933 ms）。试验阶段负载均值 3.3 时测到 174–186 / 430–454 ms，8 线程的数字对负载敏感得多。

负载高到一定程度，8 线程 engine 会突然变得非常慢。一次在负载均值从 10 涨到 97 的时段里跑的 8 线程分段 TRG（另一个 8 线程进程同时在跑）：

beta=0.30 lnZ/N=0.790559063800966 exact=0.790559070951263 err=-7.150e-9
  run ms: value first 5688.41 steady 3954.460; value+grad first 10896.22 steady 9816.194
beta=0.40 lnZ/N=0.879363060306851 exact=0.879363820774948 err=-7.605e-7
  run ms: value first 176.15 steady 174.205; value+grad first 429.61 steady 423.019
beta=0.50 lnZ/N=1.025792858455947 exact=1.025792812694918 err=+4.576e-8
  run ms: value first 4825.48 steady 4325.209; value+grad first 12670.66 steady 9804.849

同一个程序，β = 0.40 是 174 ms，前后的 β 是 3.95 s 和 4.3 s（β = 0.60 是 3.86 s），慢 20 倍以上；数值不受影响。8 个工作线程要和别的进程抢核，一个线程被换下去，整个运算就等它。所以：并行宽度不要超过空闲的核数；比较速度时用单线程做基准，8 线程的数字附上负载。上一次运行（同一程序，负载 6–8）没有出现这种跳变。

TRG 的时间花在哪

χ = 16 时，第 4 步以后每步要分解两个 256×256 的矩阵，20 步共 32 次。单独量一次 SVD（U、S、Vᵀ 都要），矩阵取 TRG 第 11 步实际的输入（β = 0.3，只有 49 个奇异值大于 1e-12）：

js = jax.jit(lambda m: jnp.linalg.svd(m, full_matrices=False))

let (u, s, vt) = x.svd()?;   // x：256×256 的占位符；三个输出都求和，免得只算奇异值

# 负载均值 7.7 到 11.8
jax jit svd 256x256 (trg): min of 30 5.373 ms
numpy svd 256x256 (trg): min of 30 5.380 ms
scipy gesdd 256x256 (trg): min of 30 5.211 ms
scipy gesvd 256x256 (trg): min of 30 24.814 ms
jax jit svd 256x256 (trg): min of 30 5.482 ms        # XLA 单线程
tenferro native svd 256x256 width=1: first 10.061 ms, min of 30 8.997 ms (checksum -17.951184)
tenferro native svd 256x256 width=8: first 111.737 ms, min of 30 5.808 ms (checksum 13.353441)
  • JAX、NumPy、SciPy 的 gesdd 都是 5.2–5.4 ms，同一类 LAPACK 分治算法；gesvd 慢 5 倍。tenferro 的 native provider（faer）单线程 9.0 ms，8 线程 5.8 ms。
  • 32 次 SVD：JAX 约 170 ms，和整个求值的 161 ms 相当（其余各步的矩阵秩更低，分解得更快）；tenferro 单线程约 290 ms，占 337 ms 的绝大部分。求值慢一倍，几乎全是 faer 和 LAPACK 在这种秩亏矩阵上的差距。随机的 256×256 矩阵上是 7.45 ms 对 11.46 ms，比例相近。
  • 两个 checksum 不同：单线程和 8 线程给出的奇异向量符号不一样（这个矩阵有 207 个零奇异值，零空间的基怎么选都对）。TRG 的结果不受影响，lnZ/N 在两种宽度下差 1e-15 量级（怎么读这一卷 的「线程数也会改末位」），但如果你逐元素比较 U、V，要先固定规范。
  • 8 线程 engine 的第一次 SVD 112 ms，是池的启动；之后 5.8 ms。计时前每个 engine 先跑一次。

收缩只占一小部分。试验阶段的 trg_step 把四个三阶张量按链 ((S1 S2) S3) S4 收缩；改成 jnp.einsum 选的两两配对 (S1 S2)(S3 S4)，单次收缩快约一倍（einsum），但整个 TRG 只快一成。同一时段交替跑两种写法，1 线程：

# 两两配对，负载均值 7.8 到 8.9
  run ms: value first 366.28 steady 318.921; value+grad first 887.65 steady 850.249
  run ms: value first 341.92 steady 326.751; value+grad first 861.74 steady 856.448
# 链式，紧接着，负载均值 8.9 到 9.9
  run ms: value first 362.02 steady 339.191; value+grad first 997.06 steady 945.638
  run ms: value first 347.06 steady 342.894; value+grad first 945.71 steady 925.000

五个 β 上，两两配对的求值 302–341 ms、求值加梯度 803–892 ms，链式 339–357 / 925–964 ms。两种写法的结果在末位不同（β = 0.30 的 lnZ/N 位模式 ...29bdf 对 ...29bdc），u 的误差到所列位数相同。本卷其他页的分段 TRG 数字都是链式写法。要追 JAX 的速度，先换 provider（CPU provider 里 Accelerate 的 SVD 512×512 单线程 32 ms，native 串行 65 ms），再谈收缩顺序。

❗IMPORTANT
动脑筋：JAX 稳态的 value 一次约 161 ms。连续调用 10 次、不等结果，派发一共要多久？之后再 block_until_ready 等全部 10 个结果，从第一次调用算起一共多久？和「每次都等」的 10 次相比呢？

同一个 jit 函数，β = 0.3，timing_pipeline.py。

动脑筋答案

# 负载均值 25
round 0: 10 calls dispatched in 0.10 ms, all ready after 1553.8 ms | 10 calls each waited 1557.1 ms
round 1: 10 calls dispatched in 0.13 ms, all ready after 1576.4 ms | 10 calls each waited 1559.7 ms
round 2: 10 calls dispatched in 0.11 ms, all ready after 1559.3 ms | 10 calls each waited 1559.8 ms
  • 10 次派发一共 0.1 ms，每次约 10 µs，比上文单次计时的 60 µs 还少：调用都排进了执行队列，Python 一侧没有等任何东西。
  • 全部就绪 1.55 s，和每次都等的 1.56 s 一样，每次约 156 ms。这台 CPU 后端把 10 次计算一个接一个地做，派发和计算的重叠只省下派发本身的几十微秒。
  • 所以不等结果的循环量不出单次的延迟，但只要最后等一次，总时间仍然对；错的只是「每次」的数字。

小结

  • JAX：稳态调用 60–85 µs 就返回，block_until_ready 之后才是 161 / 223 ms；第一次 grad 1.2 s，其中 XLA 编译 0.5 s。
  • tenferro：run 加读回就是全部时间；分段 TRG 的全部程序准备 27 ms，进程里第一次执行有几毫秒的初始化。
  • 报最小值和中位数两个数：JAX 的 jit(grad) 在负载下二者差七成，tenferro 单线程差 2%。
  • 写明 engine 宽度和负载。满载的机器上 8 线程 engine 一次可以慢 20 倍；关掉 XLA 多线程对这个 TRG 没影响。
  • TRG 的前向几乎全是 256×256 的 SVD：LAPACK gesdd 5.4 ms，faer 单线程 9.0 ms、8 线程 5.8 ms；收缩改成两两配对只快一成。

💬 关联讨论

发起讨论
加载中...