Differentiable Fortran with LFortran and Enzyme
4 hours ago
- #automatic-differentiation
- #fortran
- #scientific-computing
- Differentiable Fortran allows backpropagation through existing Fortran, C, or C++ simulation code using LFortran and Enzyme, enabling integration with modern ML frameworks like JAX and PyTorch.
- Enzyme applies autodifferentiation at the LLVM IR level, allowing differentiation of any code that compiles to LLVM, avoiding the need to rewrite validated legacy physics code.
- LFortran is used instead of Flang because it emits cleaner LLVM IR with plain pointer arithmetic, making it easier for Enzyme to trace memory accesses without opaque runtime calls.
- The compilation pipeline involves six steps: converting Fortran to LLVM IR, mild optimization, linking with a C wrapper, synthesizing derivatives via Enzyme, and final optimization into a shared library.
- Gradients from Enzyme match analytic derivatives to high precision (~6e-12 relative error), outperforming finite differences, which suffer from truncation vs. roundoff trade-offs.
- Tesseract wraps the differentiated solver as a custom JAX primitive, allowing Fortran code to be used as a differentiable layer in JAX workflows, demonstrated in a thermal forensics inverse problem.
- A case study shows reconstructing a 900-element initial temperature field from 100 noisy observations using L-BFGS-B optimization, with gradients computed via reverse-mode AD at ~450x speedup over finite differences.
- The approach is experimental and requires adapting code to avoid untraceable constructs like allocatables, using mild optimization before AD, and debugging at the IR level, but it offers composability with other AD sources.