ブラウザで密行列を速く解く — 適応周波数掃引と WebGPU の混合精度

電磁界シミュレータの計算は、ほとんどが 1 つの作業に費やされます。周波数ごとに、未知数 N の複素密行列を解くことです。手間は N³ で増えるので、N が 2 倍になれば 8 倍、10 倍になれば 1000 倍かかります。

この記事では、これをブラウザの中だけで速くするためにしたことを書きます。結論から言うと、論文の再現に使った N = 2157 の構造（21 周波数点）は、最初の約 29 秒から 2.1 秒 になりました（Apple M1）。

まず測る：JavaScript の天井

最初に、普通の JavaScript でどこまで行けるかを測りました。

行列 Z は複素対称（転置しても同じ。エルミートではない）なので、対称性を使う LDLᵀ 分解で計算量が LU の半分になります。小さなピボットが出たら、部分ピボット付きの LU に自動で切り替えます。

ブロック化が効かなかったのは、予想と原因が違ったからです。メモリ帯域が律速だと考えていましたが、実測の 3.7 GFLOP/s はメモリ転送に直すと約 15 GB/s で、M1 の帯域（約 68 GB/s）を使い切っていません。律速はベクトル命令を使えないスカラーの JavaScript の演算そのものでした。ここで JavaScript の中の細かい最適化は打ち止めにしました。

Python と numpy（Apple の Accelerate）なら 103〜123 GFLOP/s で 12〜15 倍速いことも確かめましたが、「ブラウザで開くだけ」という価値を壊すので見送りました。

メモリも問題でした。N = 8000 では行列 1 枚が 512 MB になり、作業用の複製まで含めると約 2.9 GB でタブが落ちかねません。L・Ψ・Z がすべて対称であることを使って下三角だけを詰めて持つようにし、分解も行列を上書きしながら行うことで、ピークを約 1 GB に抑えました。

周波数の点数を減らす：適応周波数掃引

次に、解く回数そのものを減らしました。

前の記事で書いたとおり、行列は

Z(ω) = jωL − jΨ/ω + Zs(ω)·G

の形をしていて、L・Ψ・G は周波数に依存しません。そこで、少数の周波数でだけ厳密に解き、その解ベクトルが張る空間に問題を縮めることにしました（縮約モデル。Momentum の Adaptive Frequency Sampling と同じ考え方です）。

1. いくつかの周波数で厳密に解き、解の実部と虚部を正規直交化して実数の基底 V を作る
2. Z_r = jω·VᵀLV − jVᵀΨV/ω + Zs·VᵀGV という小さな行列（q×q、q は 20〜60 程度）を作る。これは任意の周波数で一瞬で解ける
3. 候補の周波数すべてで残差 ‖Z(ω)·V·y − b‖ を見積もり、最も大きいところを次に厳密に解く
4. 残差が 1e-4 を下回ったら終わり

V が実数なので、縮約した行列も元の行列と同じく対称で、相反性と受動性が保たれます。厳密に解いた周波数では、厳密解と完全に一致します。残差の見積もりも、q×q の小さな行列の計算だけで済みます。

既定のプリセットでは、41 点の掃引に必要な厳密解は 8 点で、全点を厳密に解いた結果との差は 3e-13 でした。出力する点数を増やしても計算時間はほとんど変わらないので、点数の上限を 1001 に上げました。S パラメータのグラフの横軸に並ぶ小さな三角が、厳密に解いた周波数です。

倍精度の無い GPU で、倍精度の答えを得る

厳密解の回数が 10 回前後まで減ると、今度は 1 回の分解（約 3 GFLOP/s）が律速になります。ブラウザから GPU を使える WebGPU を試すと、行列積の単体で 550〜630 GFLOP/s 出ました。ただし WebGPU には倍精度（64 bit）がありません。単精度（32 bit）の分解では、有効数字が足りません。

そこで混合精度の反復改良を使いました。

1. GPU で、単精度のまま行列を LU 分解する（幅 32 のブロックに分け、行列積は 64×64 のタイルで計算）
2. その分解で連立方程式を解き、近似解 x を得る
3. 残差 r = b − Z·x を倍精度で計算する
4. 単精度の分解で補正量を解き、x に足す。2 に戻る

4〜7 回の反復で、CPU の倍精度の直接解と相対 1e-11 程度で一致します。N = 2157 の分解は GPU で 74 ms（CPU では約 4 秒）で、1 周波数ぶんの解が反復改良込みで 0.16〜0.25 秒になりました。ピボットを使わない分解なので、反復が収束しない周波数があれば、その周波数だけ自動で CPU の直接解に切り替えます。

シェーダコンパイラが誤差を消していた

手順 3 の倍精度の残差も、GPU で計算したくなります。単精度の数を 2 つ組み合わせて約 48 bit の精度を出す方法（double-float）があり、その核心は次のような「誤差のない足し算」です。

s = a + b
e = b − (s − a)    // 丸めで失われた分

ところが、Apple の GPU で動かすと、誤差の項 e が常に 0 になりました。調べると、シェーダのコンパイラが高速化のために (a + b) − a を数学的に b と簡約していました。浮動小数点では成り立たない変形で、この方法の前提を壊します。

対策は、中間結果を一度整数として扱わせることでした。外から渡した 0 と排他的論理和（XOR）を取ってから浮動小数点に戻すと、コンパイラはその値を式の木から切り離して扱い、簡約できなくなります。値そのものは変わりません。この処理を入れると、乱数ベクトルで CPU の倍精度と 1e-14 で一致するようになりました。

同じことが別の GPU で起きないとは言えないので、起動後の最初の計算で、乱数ベクトルを使って CPU と突き合わせ、合わなければ残差の計算を CPU に任せるようにしています。

前計算も GPU へ

求解が速くなると、周波数に依らない行列 L・Ψ の組み立てが目立ってきました（N = 6614 で約 7 秒）。ここも GPU に移しました。

• 遠く離れたセルどうし（全体の約 98〜99 %）は、セルの積分が点の近似で十分になり、計算が単純な級数だけになる。これを GPU で全組み合わせについて計算する
• 近くのセルどうしと対角は、CPU が元の厳密な式で計算して上書きする。近くの組は空間を区切って探し、同じ形の組は計算結果を使い回す

GPU 経路の結果と元の倍精度計算との差は、相対 1e-15 程度です。前計算は約 15 倍速くなりました。

結果

論文の再現に使った N = 2157 の構造、21 周波数点（Apple M1）：

時間が制約でなくなると、残る制約はメモリです。GPU 側に約 20·N² バイト、ブラウザ側に約 8·N² バイトが要ります。足りなくなるとエラーではなくタブが落ちうるので、未知数の上限を環境から決めています。GPU が使えれば最大 12000、使えなければ 8000 です。全面導体の N = 11402 は、41 点の掃引が 123 秒で終わりました（厳密解 14 点）。

教訓

• 推測ではなく測る。ブロック化はメモリ律速だと思って試したら、原因は演算そのものでした
• アルゴリズムで回数を減らす方が、1 回を速くするより効くことが多い。適応周波数掃引は、点数が多いほど効きます
• 精度を落とした計算は、必ず検算の仕組みとセットにする。混合精度の反復改良は残差という検算を内蔵していて、コンパイラの最適化のような思わぬ落とし穴も、自己テストで捕まえられます

次の記事では、この計算を使って論文の構造を再現しようとした話を書きます。→ ピクセル形状の回路を解く