LFortranとEnzymeによる微分可能なFortran

← Tesseract Blog 2026-07-09 · @dionhaefner LFortranとEnzymeによる微分可能なFortran 既存のFortran、C、C++のシミュレーションコードに対してバックプロパゲーションを行い、それをJAXやtorchに埋め込み、高性能な微分可能物理エンジンとして使用できたらどうなるでしょうか？ 実は、あなたが十分に勇敢であれば、それは可能です… CFD、気候、航空宇宙、原子力分野における数十年にわたる検証済みの物理コードは、現代のMLパイプラインが越えられない壁の向こうにありますが、それは勾配が公開されていないためです。通常の解決策は、すべてをJAXやPyTorchで書き直すことです。ここで探求する代替案は、コードをそのままにして、LLVMレベルの魔法のおかげで、それでも正確な勾配を得ることです。これは、EnzymeがLLVM IRレベルで自動微分を適用するため、LLVMにコンパイルされるあらゆるコードを微分できるからです！必要なのは、LFortran、LLVM、Enzymeを一緒にダクトテープで貼り付け、結果をFortranの熱ソルバーに向け、反対側から正確な勾配を得ることだけです。そこから、Tesseractは結果をカスタムJAXプリミティブとしてラップするため、Fortranソルバーは任意のJAXコード内の微分可能なレイヤーになります。これらはすべてかなり実験的であるため、NaNを返した勾配を追いかけ、手動でLLVM IRの差分を比較して機能させるために、ある程度の時間を費やす準備をしておいてください。しかし、労力を惜しまなければ、解析的な答えと一致する多段階時間ループ全体の勾配を得ることができ、最も古く最も新しい技術の組み合わせがすべて連携できるのを見るのは驚くべきことです。しかし、最初から始めましょう。以下は、コンパイルパイプライン、私たちが直面した困難な点、そして挑戦的な逆問題への応用（ADなしでは不可能だったでしょう）を含む、全ストーリーです。 レガシーコードから勾配を得る問題 科学計算に携わっているなら、この問題に遭遇したことがあるかもしれません。Fortran（またはC、C++）で書かれたシミュレーションがあり、今、誰かが勾配（シミュレーションの出力の入力に対する微分）を必要としています。最適化のためかもしれませんし、逆問題のためかもしれませんし、シミュレーションをMLパイプラインに組み込むためかもしれません。さて、あなたの選択肢は次のとおりです。 手書きのアジョイント（adjoint）。コードの各行の導関数を手動で実装することに尽きます。これは、専門家による数ヶ月の労力を意味します。エラーが発生しやすく、フォワードコードからゆっくりと同期がずれていくメンテナンスの悪夢です。 有限差分。すべての入力を摂動させ、出力を差分します。遅い（n個のパラメータに対してO(n)回の評価）、不正確で、スティッフな問題に対して条件が悪いです。 JAXまたはPyTorchで書き直し、自動微分を使用します。確かに、数万行の検証済み物理コードを書き直したいのであれば。 ソースコードから自動的に導関数をコンパイルできたらどうでしょうか？ これが実際に試してみるときの様子です。 例としてのFortran熱ソルバー 私たちのテスト対象はthermal_2d.f90で、この実験のために書いた約220行のプレーンなFortran 90です。温度依存の伝導率を持つ2次元過渡熱伝導を解きます： ρ c_p ∂T/∂t = ∇・(k(T) ∇T) + Q 伝導率kは線形材料モデル k(T) = k_0 + k_1・T に従い、時間積分はn_stepsステップでの陽的オイラー法です。ここでは何もエキゾチックなものはありません。以下は、サブルーチンシグネチャと内部ステンシルループの様子です。 ! 微分するソルバー：プレーンなFortran 90での2D熱伝導ステップループ subroutine thermal_2d_solve(n, nx, ny, n_steps, & T_init, T_final, T_cur, T_new, & k0, k1, rho, cp, & h_conv, T_inf, T_hot, & Q, Lx, Ly, dt) implicit none integer, intent(in) :: n, nx, ny, n_steps double precision, intent(in) :: T_init(n) double precision, intent(out) :: T_final(n) ! ... (ワーク配列、スカラー) ! 時間積分ループ do step = 1, n_steps do j = 2, ny - 1 do i = 2, nx - 1 idx = (j - 1) * nx + i T_c = T_cur(idx) T_e = T_cur(idx + 1) T_w = T_cur(idx - 1) T_nn = T_cur(idx + nx) T_s = T_cur(idx - nx) ! セル面での調和平均伝導率 kx_east = 2.0d0 * (k0 + k1*T_c) * (k0 + k1*T_e) & / ((k0 + k1*T_c) + (k0 + k1*T_e)) kx_west = 2.0d0 * (k0 + k1*T_c) * (k0 + k1*T_w) & / ((k0 + k1*T_c) + (k0 + k1*T_w)) ! ... (ky_north, ky_south も同様) flux_x = (kx_east*(T_e - T_c) - kx_west*(T_c - T_w)) / (dx*dx) flux_y = (ky_north*(T_nn - T_c) - ky_south*(T_c - T_s)) / (dy*dy) T_new(idx) = T_c + dt/(rho*cp) * (flux_x + flux_y + Q(idx)) end do end do ! ... (境界条件、T_cur <- T_new のスワップ) end do 境界条件は混合型です：ディリクレ（底面の熱壁）、対流/ロビン（上面）、断熱/ノイマン（側面）。面は調和平均伝導率を使用しており、これは異なるkを持つセル間のフラックス連続性のための標準的なトリックです。この設定では、導関数を得るための簡単な近道はありません。なぜなら、k(T)が非線形であるため、ステンシル係数は現在の温度場に依存し、ヤコビアンは各時間ステップで変化するからです。その手書きのアジョイントを、フォワードコードが変更されるたびに再導出することは、ADが存在する理由を削除するような作業です。 これは注釈やAD対応の構造を持たないプレーンなFortranですが、「プレーン」に到達するだけでもいくらかの調整が必要でした。LFortranでは、allocatable、array intrinsics、bounds checkingはすべてランタイムコール（_lfortran_mallocなど）にコンパイルされ、Enzymeはそれらを透過できません。そのため、ここではそれらを避ける必要があります。代わりに、ワーク配列はCから事前に割り当てられたものが渡され、--no-array-bounds-checkingでコンパイルします。220行のソルバーであれば1時間程度の作業ですが、大規模なレガシーコードベースでは未解決の問題です。 Enzymeのコンパイルパイプライン ここからは技術的な話になりますので、低レベルの詳細に興味がない場合は、次のセクションの勾配のベンチマークにスキップしてください。EnzymeはソースレベルではなくLLVM IR（中間表現）レベルで動作するため、LLVM IRにコンパイルされるもの（C、C++、Rust、Fortran）はすべて微分可能です。Fortranの場合、フロントエンドを選択する必要があります。明白な選択肢は、LLVMの公式FortranコンパイラであるFlangですが、ここではLFortranを採用します。これはEnzymeにとってより扱いやすいことが判明しました。フロントエンドが決まれば、チェーンは6回のopt/clang呼び出しで構成され、驚くほど簡単です（ただし、問題が発生した場合はIRを読んで理由を見つけることになるでしょう）。 Fortranソースから微分可能な共有ライブラリへの6段階コンパイルパイプライン。 LFortranはソルバーをLLVM IRに低下させ、薄いCラッパーがリンクされ、Enzymeが派生コードを合成し、最終的な最適化パスがフォワード、JVP、VJPエントリポイントを公開する単一の.soを出力します。チェーン全体はtesseractビルド中に約30秒で実行されます。 順序どおりの6つのステップ： # コンパイルパイプライン：Fortranソース → 微分された共有ライブラリ # 1. Fortran → LLVM IR（Enzymeがトレースするのに十分クリーン） $ lfortran --show-llvm --no-array-bounds-checking thermal_2d.f90 > thermal_2d.ll # 2. ADパスの前の軽いクリーンアップ — -O1、意図的に-O3ではない（下記参照） $ opt -O1 -S thermal_2d.ll -o thermal_2d_opt.ll # 3. CラッパーをIRにコンパイル $ clang -emit-llvm -S -O1 wrapper.c -o wrapper.ll # 4. ソルバーとllvm-linkする $ llvm-link wrapper.ll thermal_2d_opt.ll -S -o combined.ll # 5. 派生を合成する $ opt --load-pass-plugin=LLVMEnzyme-19.so -passes=enzyme combined.ll -o ad.ll # 6. 最終最適化と共有ライブラリのエミット $ opt -O3 -S ad.ll -o ad_opt.ll && clang -shared -O3 ad_opt.ll -o thermal_2d.so -lm これらのほとんどは配管作業です。実際に重要な3つのステップは、ステップ1（LFortran IRダンプ）、ステップ2（Enzyme前の-O1クリーンアップ）、ステップ5（Enzymeパス自体）であり、それぞれが検討に値する決定に対応します。 なぜFlangではなくLFortranなのか。 最初のステップ（lfortran --show-llvm）は、LFortranが驚くほどクリーンなIRを出力するため、全体が扱いやすい理由でもあります。配列は、Flang（明白なデフォルト）が生成するマルチフィールド記述子構造体やランタイムコールではなく、Cに似た通常のGEP/ロード/ストアパターンとしてプレーンなポインタとして出力されます。これは、Enzymeにとって非常に重要です。