Pangram verdict · v3.3
We believe that this document is a mix of AI-generated, AI-assisted, and human-written content
AI likelihood · overall
MixedArticle text · 1,756 words · 7 segments analyzed
← Tesseract Blog
What if you could backpropagate through existing Fortran, C, or C++ simulation code, embed it into JAX and torch, and use it as a high-performance differentiable physics engine? Turns out, you can — if you’re brave enough… Decades of validated physics code in CFD, climate, aerospace, and nuclear sit behind a wall that modern ML pipelines can’t cross, because they don’t expose gradients. The usual answer is to rewrite it all in JAX or PyTorch. The alternative we explore here is to leave the code where it is and get exact gradients out anyway, thanks to some LLVM-level magic. This is possible because Enzyme applies autodiff at the LLVM IR level, so we can differentiate any code that compiles to LLVM! All we need is to duct-tape LFortran, LLVM, and Enzyme together, point the result at a Fortran thermal solver, and get exact gradients out the other end. From there, Tesseract wraps the result as a custom JAX primitive, so a Fortran solver becomes a differentiable layer in arbitrary JAX code. This is all pretty experimental, so be prepared to spend some time chasing a gradient that returned NaN and manually compare LLVM IR diffs to make it work. But if you put in the labor, you’ll get gradients through the entire multi-step time loop that match an analytic answer, and it’s amazing to see that a stack combining some of the oldest and newest technologies can work together at all. But let’s start at the beginning. What follows is the full story, including the compilation pipeline, the sharp edges we hit, and application on a challenging inverse problem (that wouldn’t be possible without AD).
The problem with getting gradients out of legacy code¶ If you work in scientific computing, you might have run into this: you have a simulation written in Fortran (or C, or C++), and now someone needs gradients (the derivatives of the simulation’s outputs with respect to its inputs). Maybe it’s for optimization, maybe inverse problems, maybe plugging the sim into an ML pipeline. Now your options are:
Hand-written adjoints. Boils down to hand-implementing the derivative of every line of code, which means months of expert effort.
Error-prone and a maintenance nightmare that slowly drifts out of sync with the forward code. Finite differences. Perturb every input and difference the outputs. Slow (\(O(n)\) evaluations for n parameters), inaccurate, and poorly conditioned for stiff problems. Rewrite in JAX or PyTorch and use autodiff. Sure, if you want to rewrite tens of thousands of lines of validated physics.
What if you could just compile the derivatives automatically, from the existing source? This is what it actually looks like when you try.
An example Fortran heat solver¶ Our test subject is thermal_2d.f90, about 220 lines of vanilla Fortran 90 that we wrote for this experiment. It solves 2D transient heat conduction with temperature-dependent conductivity:
\[\rho \, c_p \frac{\partial T}{\partial t} = \nabla \cdot \big( k(T) \, \nabla T \big) + Q\]
The conductivity \(k\) follows a linear material model, \(k(T) = k_0 + k_1 \cdot T\), and time integration is explicit Euler over n_steps steps. Nothing exotic here. Here’s what the subroutine signature and the interior stencil loop look like: ! The solver we differentiate: a 2D heat conduction step loop in plain Fortran 90 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) ! ... (
work arrays, scalars)
! Time integration loop 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)
! Harmonic-mean conductivity at cell faces 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 similarly)
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 ! ... (
boundary conditions, swap T_cur <- T_new) end do
Boundary conditions are mixed: Dirichlet (hot wall at the bottom), convection/Robin (top), and insulated/Neumann (sides). The faces use harmonic-mean conductivity, the standard trick for flux continuity across cells with different \(k\). For this setup, there are no easy shortcuts to get derivatives: because \(k(T)\) is nonlinear, the stencil coefficients depend on the current temperature field, so the Jacobian changes at every time step. Hand-coding the adjoint through that, and re-deriving it every time the forward code changes, is the kind of work AD exists to delete. This is vanilla Fortran with no annotations or AD-aware constructs, but even getting to “vanilla” took some massaging. In LFortran, allocatables, array intrinsics and bounds checking all compile down to runtime calls (_lfortran_malloc and friends) that Enzyme can’t see through, so we have to avoid them here. Instead, work arrays get passed in pre-allocated from C, and we compile with --no-array-bounds-checking. For a 220-line solver that’s an hour of work, while for a large legacy codebase it’s an open question.
The Enzyme compilation pipeline¶ This is where it gets technical, so if you’re not into the low-level details, feel free to skip to the next section where we benchmark the gradients. Enzyme works at the LLVM IR (intermediate representation) level, not the source level, so anything that compiles to LLVM IR (C, C++, Rust, Fortran) can be differentiated. For Fortran that means picking a frontend. The obvious choice is Flang, LLVM’s official Fortran compiler, but we reach for LFortran instead, which turns out to be much more digestible for Enzyme. With a frontend settled, the chain consists of six opt/clang invocations and is surprisingly straightforward (though when things go wrong you’ll end up reading the IR to find out why).
The six-step compilation pipeline, from Fortran source to a differentiable shared library. LFortran lowers the solver to LLVM IR, a thin C wrapper is linked in, Enzyme synthesizes the derivative code, and a final optimization pass emits a single .so exposing forward, JVP, and VJP entry points.
The whole chain runs in about 30 seconds during tesseract build.¶
The six steps, in order: # The compilation pipeline: Fortran source → differentiated shared library
# 1. Fortran → LLVM IR (clean enough for Enzyme to trace) $ lfortran --show-llvm --no-array-bounds-checking thermal_2d.f90 > thermal_2d.ll
# 2. Mild cleanup BEFORE the AD pass — -O1, deliberately not -O3 (see below) $ opt -O1 -S thermal_2d.ll -o thermal_2d_opt.ll
# 3. Compile the C wrapper to IR $ clang -emit-llvm -S -O1 wrapper.c -o wrapper.ll
# 4. llvm-link it with the solver $ llvm-link wrapper.ll thermal_2d_opt.ll -S -o combined.ll
# 5. Synthesize the derivative $ opt --load-pass-plugin=LLVMEnzyme-19.so -passes=enzyme combined.ll -o ad.ll
# 6. Final optimization and emit a shared library $ opt -O3 -S ad.ll -o ad_opt.ll && clang -shared -O3 ad_opt.ll -o thermal_2d.so -lm
Most of those are plumbing. The three that actually matter are step 1 (the LFortran IR dump), step 2 (the pre-Enzyme -O1 cleanup), and step 5 (the Enzyme pass itself), where each maps to a decision worth dwelling on.
Why LFortran, not Flang. The first step (lfortran --show-llvm) is also the reason the whole thing is tractable, since LFortran emits remarkably clean IR. Arrays come out as plain pointers with the usual GEP/load/store patterns, much like C, instead of the multi-field descriptor structs and runtime calls that Flang, the obvious default, produces. That matters enormously, because Enzyme has to trace every memory access to figure out what’s active, and it has a far easier time with pointer arithmetic than with opaque fir.box descriptors. So what makes LFortran IR “clean”? Three things, concretely. Every array reference lowers to a plain getelementptr + load/store on a bare pointer; the 19-argument subroutine takes 19 plain ptr arguments, with no descriptor struct wrapping any of them; and the whole module contains zero function calls: no _lfortran_malloc, no runtime helpers, nothing opaque for Enzyme to trace through. Here’s the IR LFortran emits for the start of the stencil body, covering the idx = (j-1)*nx + i index and the T_c = T_cur(idx) / kx_east = ... lines from the source above: ; What LFortran emits for the stencil body: plain pointer math, no descriptors ; idx = (j - 1) * nx + i %56 = sub i32 %55, 1 ; j - 1 %58 = mul i32 %56, %57 ; (j-1) * nx %60 = add i32 %58, %59 ; + i store i32 %60, ptr %idx
; T_c = T_cur(idx) %67 = getelementptr inbounds
double, ptr %t_cur, i32 %65 ; &T_cur[idx] %68 = load double, ptr %67 ; T_cur[idx] store double %68, ptr %t_c
; kx_east = 2*(k0+k1*T_c)*(k0+k1*T_e) / ((k0+k1*T_c)+(k0+k1*T_e)) %110 = fmul double %108, %109 ; k1 * T_c %111 = fadd double %107, %110 ; k0 + k1*T_c %112 = fmul double 2.0, %111 ; ... k0 + k1*T_e, the product, the sum ... %130 = fdiv double %118, %129 ; the harmonic mean store double %130, ptr %kx_east
That’s it. Array access is a getelementptr + load, arithmetic is a flat sequence of fmul/fadd/fdiv that maps line-for-line onto the Fortran. Flang, by contrast, would wrap those same arrays in descriptor structs and reach for runtime helpers, and Enzyme would have to see through every one of them. The catch is that LFortran is still maturing, so your code has to stay inside what it supports. Why -O1, not -O3. Our first pipeline used -O3 here, and it cost us hours: the forward pass was perfect, but the VJP returned NaN on certain inputs. At -O3, LLVM’s aggressive vectorization and code motion produce IR patterns Enzyme seems to mishandle in reverse mode. In our case it bit when adjacent cell temperatures were equal, so intermediate terms cancelled and a compiler rearrangement turned that into a division by zero. The fix is to keep optimization mild before the AD pass and save -O3 for after it. C bridge. In between Fortran and Enzyme sits a thin C wrapper that bridges Fortran’s by-pointer ABI to a C interface Enzyme can annotate.