Differentiable Fortran with LFortran and Enzyme

背景

科学计算领域积累了数十年的验证代码——CFD、气候、航空航天、核工程——这些代码大多用 Fortran、C 或 C++ 写成，但其输出相对于输入的导数（梯度）却无法直接获取。现代机器学习流水线（如 JAX、PyTorch）需要可微分的组件，而传统仿真代码暴露不出梯度，导致这些经过严格验证的物理求解器无法嵌入到 AI 管道中。

常规做法是将代码全部重写成 JAX 或 PyTorch，但对于数万行的 Fortran 代码，重写是不现实的。另一种思路是保持原代码不动，借助底层编译器技术自动提取精确梯度。Enzyme 是一个在 LLVM IR（中间表示）层面进行自动微分的工具，任何能编译到 LLVM 的语言（C、C++、Rust、Fortran）都可以用它来求导。LFortran 是一个 Fortran 编译器，它生成的 LLVM IR 特别干净，易于 Enzyme 处理。将 LFortran、LLVM 和 Enzyme 组合起来，就能让现有 Fortran 热求解器输出精确梯度，并通过 Tesseract 包装成 JAX 自定义原语，使 Fortran 求解器变成 JAX 中可微分的层。

核心内容

从遗留代码获取梯度的困境

如果你从事科学计算，可能遇到过这样的场景：有一个用 Fortran（或 C/C++）编写的仿真程序，现在需要输出的梯度（关于输入的导数）。用于优化、反问题或接入 ML 管道。通常的选项包括：

• 手写伴随代码：需要为每行代码手动实现导数，花费数月专家精力，容易出错且维护噩梦（会逐渐与前向代码失同步）。
• 有限差分法：扰动每个输入并差分输出。对于 n 个参数需要 O(n) 次评估，速度慢、不精确，且对刚性条件病态。
• 重写在 JAX/PyTorch 中：需要重写数万行验证过的物理代码。

我们希望直接从现有源代码自动编译出导数。

示例：Fortran 热求解器

测试对象是 thermal_2d.f90，约 220 行标准 Fortran 90，用于求解二维瞬态热传导，热导率随温度线性变化（k(T) = k₀ + k₁·T），时间积分为显式欧拉。子例程签名与内部模版循环如下：

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)

    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 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
    end do

边界条件混合：底部 Dirichlet（热壁）、顶部对流/Robin、侧面绝热/Neumann。面处使用调和平均热导率。由于 k(T) 非线性，模板系数依赖当前温度场，雅可比矩阵每一步都在变化，手工编码伴随体非常繁琐，这正是自动微分要消除的工作。

为了能让 LFortran 编译，需要对原代码做一些调整：避免使用可分配数组、数组内建函数和边界检查（这些会编译成运行时调用 _lfortran_malloc 等，Enzyme 无法穿透）。工作数组预先由 C 分配并传入，编译时使用 --no-array-bounds-checking。对于一个 220 行的求解器，这是一小时的工作；对大型遗留代码库则是一个开放问题。

Enzyme 编译管道

Enzyme 在 LLVM IR 级别工作，因此任何能编译到 LLVM IR 的代码都可微分。Fortran 的前端可选 Flang 或 LFortran。LFortran 因其生成的 IR 更干净而被选用。整个管道包含六步 opt/clang 调用：

# 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 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

关键步骤：

• 第 1 步：lfortran --show-llvm 是让整个方法可行的关键。LFortran 生成非常干净的 IR：数组变成纯指针和 GEP/load/store 模式，类似于 C，而不是 Flang 产生的多字段描述符结构体和运行时调用。这使得 Enzyme 可以轻松追踪每个内存访问以确定哪些变量是活跃的。
• 第 2 步：使用 -O1 而不是 -O3，因为过高的优化可能会重排代码，使 Enzyme 难以正确理解依赖关系。
• 第 5 步：Enzyme 传递本身，在整个组合 IR 上应用自动微分。

为什么 LFortran IR 干净

具体来说有三点：

1. 每个数组引用都简化为直接的 getelementptr + load/store，基于裸指针。
2. 19 个参数的子例程接受 19 个裸 ptr 参数，没有描述符结构体包裹其中任何一个。
3. 整个模块包含零个函数调用：没有 _lfortran_malloc，没有运行时调用。

关键要点

• Enzyme 在 LLVM IR 层面做自动微分，因此任何能编译到 LLVM 的语言（Fortran、C、C++、Rust）都可以被微分，无需修改源代码。
• LFortran 而非 Flang