FPU エミュレータを超最適化する

超最適化（superoptimization）とは、コンパイラがループを含まない命令列に対して最適な命令列を自動的に見つけ出すことだ。

– Wikipedia

コンパイラは機械的な変換でコードを速くするのは得意だが、大半のコンパイラは自明でない推論は苦手で、目の前の問題を解く新しい発想を思いつくことはできない。

浮動小数点演算ライブラリを設計していたとき、浮動小数点数 $f$ を受け取り、$f$ のうち最も下にあるセットビットに対応する浮動小数点数を返す関数が必要になった。たとえば $f=1.01101\cdot 2^{10}$（2進数）なら $0.00001\cdot 2^{10}=2^5$ を返す。

ビット演算で実装できるが、最も素直なアルゴリズムは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 は見落とすし、よくあるパターンに偏りすぎる。総当たりならそれはない。

だが、式が正しいかどうかをどう確かめるのか。定理証明器は遅すぎる。代わりに代表的なテストケースをいくつか用意し、見つかった式が正しいことを祈る。反例が出たらその反例をテストケースのリストに加えて、もう一度やり直す。

それほど無茶な話でもない。この時点で私は $6$ 個の演算を使うアルゴリズムを実装していたので、演算数が $5$ 個以下の式を列挙すればよい。次の演算が役に立ちそうだというのは分かっていた：

• ビットトリック用の整数の加減算、およびビット単位の AND/OR/XOR。
• 最後の演算としての浮動小数点の減算。
• 比較も要るかもしれない。

$5$ 個の演算を選び、それぞれが $6$ 種類の変種のどれかになりうるので $6^5=7776$ 通り、さらに $C_5=42$ 通りの括弧の付け方を掛ける。大した数ではない！

ただし、各演算の入力も選ばなければならない。入力は入力値 $x$ か、何らかの定数だ。定数は完全にランダムにする必要はなく、float の2進表現の各フィールドを覆うものだけで足りると見込んだので、定数のリストは $64$ 個まで減る。組み合わせは ${65}^6$ 通りになり、多いように聞こえるが、実際にはかなり絞り込める。

探索

探索を効率化するため、帰納法で式の集合を段階的に構築した。

$0$ 個の演算を含む式は $x$ と定数だけ。$n$ 個の演算を含む式は、$k$ 個と $n-1-k$ 個の演算を含む式を組み合わせることで得られる：

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])
                        }),
                    });
                }
            }
        }
    }
}

帰納法のおかげで重複を除ける。たとえば $(x+1)+2$ と $(x+2)+1$ は別の式だが、返す値は同じだ。出力ごとに式を1つだけ残せば、以降の反復は速く終わる。おかげで総当たり探索は $4$ 分で走り切った！（まだ最適化の余地はあるはずだが、一度しか使わないスクリプトを最適化するくらいなら YouTube でも見ていたい。）

時間とメモリを節約するため、最後の反復では重複除去をせず、一致する候補をすべてそのまま出力する。メモリが問題になるのは、この時点でテストケースが $20$ 個あり、式1つにつき $80$ バイトを保存する必要があり、しかも式の数が非常に多いからだ。

結果

この総当たり探索は何か面白いものを見つけたのか。見つけた！ $5$ 個の演算を含む式がいくつかあり、たとえばこれだ：

return x - f32::from_bits(x.to_bits() & (
    (f32::EXPONENT_MASK - (x | f32::EXPONENT_MASK)) ^ !f32::EXPONENT_MASK
));

$4$ 個の演算を含む式は見つからなかったが、この手法には限界がある。計算の再利用を考慮していないのだ。たとえば x & 1 が2か所で使われても、2個の演算として数えられ、1個とは数えられない。$5$ 個の演算の式のどれかが、再利用によって実質 $4$ 個の演算で書けることを期待していた。すべての式を手で確認したところ、最後の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) に書き換え、それを反転と統合することで演算を1回節約している！

これほど巧妙な書き換えを見つけたと気づいたとき、私は度肝を抜かれた。LLM がなくても、コンピュータの発明力には驚かされる！ しかもそれが業界を破壊することもない。

ともあれ、このコードは今まだ開発中の私の FPU シミュレータの中に住んでいる。私はこれを誇りに思っている。