Legacy Fortran atmospheric physics moves to GPU: four LLMs each wrote a JAX version, and the fastest ran about 357 times faster than serial Fortran
Synopsis
Four LLMs rewrote the Kessler warm-rain microphysics Fortran code into GPU-executable JAX, reproducing the Fortran reference to near floating-point precision while differing widely in GPU performance, with the fastest reaching about 357 times the speed of serial Fortran in cached GPU execution.
Interpretation
LAFT converts Fortran atmospheric physics parameterizations into JAX that runs on GPUs while staying numerically close to the original Fortran reference. These physics packages were written for CPUs, and manual porting to GPUs requires expertise in accelerator programming, memory layout, and performance tuning, which bottlenecks modeling groups modernizing their code. On the Kessler warm-rain microphysics scheme, the maximum absolute difference between the four LLM-generated JAX implementations and the Fortran reference stayed inside the reference's own variability envelope from 22 Fortran build configurations, with ratios from 0.42 to 0.87.
Numerically equivalent translations do not run equally fast on GPUs, because the generated code exposes different computation structures to the compiler. Prior code-translation work reported compilation, runtime, and semantic errors, but rarely separated numerical correctness from execution efficiency. On a single NVIDIA A100, the Claude and GPT versions took 7.32 ms and 7.26 ms at 10000 columns, the Gemini version took 3.05 ms, and serial Fortran took 496.194 ms, while the Qwen version kept a sequential column loop and took 6271.75 ms at 10000 columns.
LAFT separates memory-layout conversion from the scientific computation, using a deterministic front end to generate a bridge layer so the LLM only writes the compute kernel. Fortran uses column-major arrays while JAX defaults to row-major, and delegating that conversion to the LLM inside every kernel is error-prone and repetitive. The bridge is validated separately for round-trip layout fidelity, transposition, shapes, and memory contiguity; at 10000 columns the input and output permutations together cost about 10.5 ms, while host-to-device and device-to-host transfers together cost about 6.75 s.
Perspective
The report covers a single call of the Kessler warm-rain microphysics scheme, with inputs generated from a fixed seed, 56 vertical levels, and column counts from 50 to 1000000. Numerical acceptance uses the reference's own variability envelope from 22 Fortran build configurations as its threshold, and performance is measured as cached execution time on a single NVIDIA A100, excluding compilation and initial host-device transfers. It is aimed at researchers and engineers working on GPU modernization of atmospheric models and scientific code translation.
The acceptance envelope is specific to one input set, one call of the scheme, and one machine, and the compilation ensemble characterizes numerical variability for the tested inputs rather than coverage of branches triggered by special values. The Qwen version retained a sequential column loop after five consecutive repair attempts and is kept to characterize its execution design. Bridge overhead is estimated by summing separately measured phases rather than measured directly in the production call. The OpenACC comparison shows OpenACC faster at small column counts and the Gemini-generated JAX faster at large column counts, so the advantage depends on problem size.