EinsumOptimize::False 反而慢几百倍
主线位置:einsum 线的第 2 步。traced 入口的
einsum_with收一个EinsumOptimize,决定收缩顺序;False的本意是「不优化,照写的顺序算」。几种选项和它们在运行时怎么变成一棵树,见 einsum。
flowchart TB l0["列表 0 1 2 3 4"] --> s1["收缩 0 1 得 5"] s1 --> l1["列表 2 3 4 5"] l1 --> s2["收缩 2 3 得 6"] s2 --> l2["列表 4 5 6"] l2 --> s3["收缩 4 5 外积"] s3 --> big["256 立方个元素"]
实验 crate exp-tenferro8/pit1-pit,bin p11_false,release;f64,traced 程序,CpuEngine::serial(),prepare 之后第一次 run 单独计,之后 20 次取最小值;跑了两轮,下面引第二轮。uptime 1 分钟负载 2.3 到 2.5。M2 Max、macOS 27、rustc 1.99.0。
现象
四个 256×256 的矩阵乘一个长 256 的向量,"ij,jk,kl,lm,m->i"。想「先别让规划器插手,按写的顺序算」,用 EinsumOptimize::False;对照默认规划和一条手写的从左到右路径:
let out = trace.einsum_with(&vals, notation, EinsumOptimize::False)?;
// 对照
let out = trace.einsum_with(&vals, notation, EinsumOptimize::default())?;
let out = trace.einsum_with(&vals, notation, EinsumOptimize::Path(vec![(0, 1), (0, 3), (0, 2), (0, 1)]))?;
== ij,jk,kl,lm,m->i, d=256, f64, serial engine
default pairs=[(3, 4), (2, 5), (1, 6), (0, 7)] flops=2.621e5 peak=2.560e2 first=0.17ms min-of-20=0.057ms y[0]=4.252080e2
False pairs=[(0, 1), (2, 3), (4, 5), (6, 7)] flops=6.711e7 peak=1.678e7 first=54.86ms min-of-20=46.505ms y[0]=4.252080e2
Path left-to-right pairs=[(0, 1), (2, 5), (3, 6), (4, 7)] flops=5.040e7 peak=6.554e4 first=2.25ms min-of-20=2.090ms y[0]=4.252080e2
== ij,jk,kl,lm,mn,no,o->i, d=256, f64, serial engine
default pairs=[(5, 6), (4, 7), (3, 8), (2, 9), (1, 10), (0, 11)] flops=3.932e5 peak=2.560e2 first=0.13ms min-of-20=0.077ms y[0]=-1.242054e5
False pairs=[(0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11)] flops=1.007e8 peak=1.678e7 first=27.53ms min-of-20=20.048ms y[0]=-1.242054e5
Path left-to-right pairs=[(0, 1), (2, 7), (3, 8), (4, 9), (5, 10), (6, 11)] flops=8.395e7 peak=6.554e4 first=3.54ms min-of-20=3.557ms y[0]=-1.242054e5
(pairs 是 ContractionTree 的步骤:引用原输入 0..n 或之前步骤的结果 n + 步号;flops 是每步两个操作数所有标签长度之积的和,peak 是最大的中间结果元素数。False 这一行的树由实验程序按它的实现规则用 ContractionTree::from_pairs 建出,计时用的是真正的 EinsumOptimize::False。)
五个输入时,False 比真正的从左到右慢 22 倍,比默认规划慢 800 多倍;七个输入时比默认慢 260 倍。最大的中间结果有 1678 万个 f64,134 MB。结果三者一致,没有任何提示。
出处
EinsumOptimize::False 的文档注释是「No optimization: contract operands left-to-right」。运行时把它解析成树的地方在 resolve_plan_spec:
EinsumPlanSpec::LeftToRight => {
let n = subscripts.inputs.len();
if n <= 1 {
ContractionTree::from_pairs(subscripts, shapes, &[])
} else {
let path: Vec<(usize, usize)> = (0..n - 1).map(|_| (0, 1)).collect();
let pairs = jax_path_to_v1_pairs(&path, n)?;
ContractionTree::from_pairs(subscripts, shapes, &pairs)
}
}
(0, 1) 重复 n - 1 次,按 JAX 路径的语义解释(jax_path_to_v1_pairs:每步收缩当前列表里的这两个位置,两者移除,结果追加到末尾)。五个输入时列表依次是 [0,1,2,3,4] → [2,3,4,5] → [4,5,6] → [6,7],四步分别收缩 (0,1)、(2,3)、(4,5)、(6,7):成对地往前推,而不是把累积的结果和下一个输入收缩。第三步的 4 是向量 m,5 是 ij,jk 的结果 ik,两者没有共享标签,收缩成外积 [m, i, k],256³ 个元素。
三个输入时 (0,1)、(2,3) 恰好就是从左到右,看不出差别。检查树形状的两个上游单元测试 false_strategy_builds_left_to_right_tree 和紧接着的 plan_spec_from_false_preserves_left_to_right_policy 都用 "ij,jk,kl->il",断言的正是 (0, 1)、(2, 3);集成测试里用到 False 的只有两个输入的 "ij,jk->ik"。从四个输入起顺序就不是从左到右了;会不会造出外积取决于网络,两两配对的两个操作数没有共享标签时就会。实现和文档不一致,是上游 bug,0.8.0 上复现,这段代码和 0.7.1 相同。
问: 路径格式是从 JAX 借的,(0, 1) 重复不就是「每次拿前两个」,也就是从左到右吗?
答: 「每次拿前两个」只有在结果放回最前面时才等于从左到右。这种路径格式(EinsumOptimize::Path 的文档注释写明「After each contraction, the two operands are removed and the result is appended」)把结果追加到末尾,所以下一步的「前两个」是还没碰过的原输入。要用这种格式表达从左到右,后续每一步都得把末尾的累积结果和当前第一个输入配对:[(0,1), (0,k), ..],k 是当前列表的最后一个位置,也就是上面的 Path(vec![(0, 1), (0, 3), (0, 2), (0, 1)])。
修法
要从左到右,就自己写 Path:n 个输入时是 (0, 1) 之后接 (0, n-2)、(0, n-3)、……、(0, 1)。
fn left_to_right_path(n: usize) -> Vec<(usize, usize)> {
let mut path = vec![(0, 1)];
for live in (2..n).rev() {
path.push((0, live - 1));
}
path
}
(n = 5 时得到 [(0, 1), (0, 3), (0, 2), (0, 1)],就是上面计时用的那条。)不过多数时候根本不需要「不优化」:默认规划(tenferro 自己的贪心,见 einsum)在这条链上找到了「从向量开始」,比从左到右快 40 倍左右。需要确定的顺序(可重复、或者手里有更好的树)时用 Path / Tree / Nested,不要用 False。
怎么确认没中招
- 用
EinsumOptimize::False的 traced 程序第一次运行特别慢、内存突然上涨:按上面的规则手推一遍位置列表,看两两配对的操作数有没有共享标签。 - 用
ContractionTree::from_pairs建出同样的树,把step_subscripts(i)打出来:某一步的两个输入标签没有交集,就是外积。
小结
EinsumOptimize::False的实现是 JAX 路径(0, 1)重复n - 1次;结果追加到末尾,从第四个输入起变成两两配对。- 矩阵链乘向量时第三步是向量和矩阵积的外积,256³ 个元素;五个输入时比默认规划慢 800 多倍,比真正的从左到右慢 22 倍。
- 文档和实现不一致,上游 bug(0.8.0 复现);要从左到右请写
Path,通常直接用默认规划更好。