超优化一个 FPU 模拟器
用暴力搜索遍历所有短的位运算公式,找到一个无分支取浮点数最低置位的写法,还发现一处复用子表达式的改写,这是编译器不会自己想到的。
中文
复制

超优化(superoptimization)指的是编译器自动为一段无循环的指令序列找出最优指令序列的过程。
编译器擅长做机械变换,让代码跑得更快,可惜大多数编译器并不擅长非平凡的推理,也就想不出解决你问题的全新思路。
在设计一个浮点计算库时,我需要一个函数,它接收一个浮点数 ,返回一个浮点数,对应 中最低的那个置位。例如,若 (二进制),它会返回 。
这可以用位运算实现,但最直接的算法在 2 的幂上就不成立了,所以我不得不插入一个 if:
if f32::from_bits(x.to_bits() & !f32::MANTISSA_MASK) == x { // is x a power of two?
return x; // x itself is the lowest bit
} else {
return x - f32::from_bits(x.to_bits() & (x.to_bits() - 1)); // x, minus x without lowest bit
}
这个 if 让我很恼火。常见的输入被它拖慢了,分支还可能预测失败。我想找一个能自动处理 2 的幂的公式,但怎么试都找不到一个足够有说服力的最优解。
于是我走了一步不太常规的棋。
暴力搜索
如果有办法自动找出这样的公式呢?不是靠召唤一群 LLM,而是暴力枚举所有可能的公式。LLM 会漏东西,也会过度偏向常见模式;暴力搜索不会。
但怎么检验一个公式是否正确?定理证明器太慢,不过我们可以设计几组有代表性的测试用例,然后指望找到的公式是对的。万一出现反例,就把反例加进测试用例列表,再来一遍。
这听起来没那么疯狂。到这一步,我已经实现了一个用 条运算的算法,所以只需要枚举运算数不超过 条的公式。我知道下面这些运算可能有用:
- 整数加减,以及按位 AND/OR/XOR,用于位技巧。
- 浮点减法,作为最后一步运算。
- 也许还有比较?
选出 条运算,每条可以是 种变体之一,一共 种,再乘上 种可能的括号位置。这也没多少!
不过,我们还得为每条运算选择输入,输入可以是输入值 ,也可以是某个常量。我猜只需要覆盖 float 二进制表示各个字段的常量就够了,不必用完全随机的常量,这样常量列表就缩减到 个。组合数是 ,听起来很多,但实际上可以筛掉不少。
搜索
为了让搜索高效,我用归纳法逐步构建公式集合。
含 条运算的公式只有 和常量;含 条运算的公式,可以通过把含 条和 条运算的公式两两组合得到:
for n_ops in 1..N_OPS {
let mut new_exprs = FxHashSet::default();
for a_ops in 0..n_ops {
let b_ops = n_ops - 1 - a_ops;
let a_exprs = &exprs_per_op_count[a_ops];
let b_exprs = &exprs_per_op_count[b_ops];
for a_expr in a_exprs {
for b_expr in b_exprs {
for op in OPS {
// New expression: `a_expr op b_expr`
new_exprs.insert(Expr {
values: core::array::from_fn(|i| {
op(a_expr.values[i], b_expr.values[i])
}),
});
}
}
}
}
}
归纳法让我们能去掉重复项。比如 和 是两个不同的表达式,但返回值相同。每组输出只保留一个表达式,后续迭代就能更快完成。正因如此,暴力搜索只要 分钟就能跑完!(我相信它还能继续优化,但我宁愿去看 YouTube,也不想优化一个只用一次的脚本。)
为了省时间和内存,最后一轮迭代不做去重,直接输出所有匹配的候选。内存之所以要紧,是因为到这一步我已经有 个测试用例,每个表达式得存 字节,而表达式的数量非常多。
结果
那么,这次穷举搜索有没有找到什么有意思的东西?找到了!有几个含 条运算的公式,比如这个:
return x - f32::from_bits(x.to_bits() & (
(f32::EXPONENT_MASK - (x | f32::EXPONENT_MASK)) ^ !f32::EXPONENT_MASK
));
我没能找到含 条运算的公式,但这个方法有个局限:它没有考虑计算可以复用。比如 x & 1 在两处用到,它会被算作两条运算,而不是一条。我原本希望某个 运算的公式能靠复用真正写成 条运算,而手动检查完每一个公式后,最后一个恰好就有这个性质!它长这样:
let a = x.to_bits() & !f32::MANTISSA_MASK; // x without mantissa bits
return f32::from_bits(x.to_bits() & (a - x.to_bits())) - f32::from_bits(a);
它在 2 的幂上表现略有不同,会翻转符号位,这一点我不在乎。这套逻辑乍看之下讲不通,但它确实管用。
只有这个公式能复用 a,靠的是子表达式 a - x.to_bits()。它等价于 -(x & f32::MANTISSA_MASK),其他公式正是这样引用它的。所有公式都会计算 x & !f32::MANTISSA_MASK 和 x & f32::MANTISSA_MASK,但这个公式把后者改写成 x - (x & f32::MANTISSA_MASK),再与取反合并,省下一次运算!
意识到它找到了如此巧妙的改写时,我惊呆了。你不需要 LLM 也能被计算机的发明能力震撼到!而且它不会摧毁任何行业。
总之,这段代码现在住在我那个仍在开发中的 FPU 模拟器里,我为此感到自豪。