LFortranとEnzymeで既存のFortranコードに自動微分を導入する方法

Differentiable Fortran with LFortran and Enzyme

CFDや気候、航空宇宙、原子力などで数十年検証されてきたFortranコードは、勾配を公開していないため現代のMLパイプラインに組み込めない。通常はJAXやPyTorchで書き直すが、この記事ではLFortran、LLVM、Enzymeを組み合わせ、既存コードを変更せずにLLVM IRレベルで自動微分を実現する方法を紹介する。220行のFortran熱ソルバーを例に、6ステップのコンパイルパイプライン、LFortranが生成するクリーンなIR、Enzymeの逆モードでのNaN問題を回避するための-O1最適化の重要性、Cブリッジによるメモリ管理などを解説。最終的にはJAXのカスタムプリミティブとして統合し、逆問題を解く。