使用LFortran和Enzyme实现可微Fortran

# 使用 LFortran 和 Enzyme 实现可微分的 Fortran 代码 来源：https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html ← Tesseract 博客 (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/index.html)如果你能通过现有的 Fortran、C 或 C\+\+ 仿真代码进行反向传播，将其嵌入 JAX 和 torch，并用作高性能可微分的物理引擎，那会怎样？事实证明，如果你够大胆，你是可以做到的…… 数十年经过验证的物理代码——CFD、气候、航空航天、核工程——都被一堵现代机器学习管道无法跨越的墙隔开了，因为它们不提供梯度。通常的答案是全部用 JAX 或 PyTorch 重写。而我们在这里探索的替代方案是，保持代码原封不动，仍能获取精确梯度，这要归功于一些 LLVM 级别的魔法。这是可行的，因为Enzyme (https://enzyme.mit.edu/)在 LLVM IR 级别应用自动微分，因此我们可以对任何能编译为 LLVM 的代码进行微分！ 我们需要的只是将LFortran (https://lfortran.org/)、LLVM 和 Enzyme 用“胶带”粘在一起，然后指向一个 Fortran 热求解器，就能从另一端得到精确的梯度。接着，Tesseract (https://github.com/pasteurlabs/tesseract-core)将结果封装成一个自定义的JAX (https://jax.readthedocs.io/en/latest/)原语，这样 Fortran 求解器就成了任意 JAX 代码中的一个可微分层。 这一切都还非常实验性，所以准备好花些时间追踪返回 NaN 的梯度，并手动比较 LLVM IR 的差异来使其工作。但如果你投入了这些劳动，你将得到贯穿整个多步时间循环的梯度，并且与解析解匹配。看到这样一个结合了最古老和最前沿技术的栈居然能协同工作，真是令人惊叹。 但让我们从头开始。接下来是完整的故事，包括编译管道、我们遇到的尖锐问题，以及在一个具有挑战性的逆问题（如果没有自动微分根本无法解决）上的应用。 ## 从遗留代码中获取梯度的问题¶ (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#the-problem-with-getting-gradients-out-of-legacy-code) 如果你从事科学计算，你可能遇到过这种情况：你有一个用 Fortran（或 C、C\+\+）编写的仿真程序，现在有人需要梯度（仿真输出相对于输入的导数）。可能是为了优化、逆问题，或者将仿真接入机器学习管道。现在你的选择有： - **手写伴随。** 这相当于手写实现每一行代码的导数，意味着需要专家投入数月时间。容易出错且维护困难，会慢慢与正向代码不同步。 - **有限差分。** 扰动每个输入并对输出求差分。速度慢（对于 n 个参数需要 \\(O(n)\\) 次评估），不精确，且对于刚性问题条件不佳。 - **用 JAX 或 PyTorch 重写并使用自动微分。** 当然，前提是你想重写数万行经过验证的物理代码。 如果你能直接从现有源码自动编译出导数呢？下面就是尝试时的实际样子。 ## 一个 Fortran 热求解器示例¶ (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#an-example-fortran-heat-solver) 我们的测试对象是 `thermal_2d.f90`，大约 220 行普通 Fortran 90 代码，是我们为这个实验编写的。它求解二维瞬态热传导，具有温度依赖的导热系数： \\\[\\rho \\, c\_p \\frac\{\\partial T\}\{\\partial t\} = \\nabla \\cdot \\big\( k\(T\) \\, \\nabla T \\big\) \+ Q\\\] 导热系数 \\(k\\) 遵循线性材料模型，\\(k\(T\) = k\_0 + k\_1 \\cdot T\\)，时间积分是显式欧拉法，进行 `n_steps` 步。没有什么特别的。 下面是子程序签名和内部模板循环的样子： `` ! 我们微分的求解器：普通 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) ! ... (工作数组，标量) ! 时间积分循环 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 类似) 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 ! ... (边界条件，交换 T_cur <- T_new) end do `` 边界条件是混合的：狄利克雷（底部热壁）、对流/罗宾（顶部）和绝缘/诺伊曼（侧面）。各面使用调和平均导热系数，这是跨不同 \\(k\\) 的网格单元之间保持通量连续性的标准技巧。 对于这个设置，没有捷径可以获取导数：因为 \\(k\(T\)\\) 是非线性的，模板系数依赖于当前温度场，因此雅可比矩阵在每个时间步都会变化。手动编写伴随代码，并且每次正向代码变化时都要重新推导，这正是自动微分所要消除的工作。 这是普通的 Fortran，没有任何注解或自动微分相关的结构，但即使要达到“普通”也需要一些调整。在 LFortran 中，可分配数组、数组内在函数和边界检查都会编译为运行时调用（`_lfortran_malloc` 等），Enzyme 无法看透它们，因此我们必须避免使用它们。相反，工作数组从 C 预分配后传入，并且我们使用 `--no-array-bounds-checking` 编译。对于一个 220 行的求解器，这需要一小时的工作，而对于大型遗留代码库，这还是一个开放问题。 ## Enzyme 编译管道¶ (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#the-enzyme-compilation-pipeline) 这部分技术性较强，如果你对底层细节不感兴趣，可以直接跳到下一节 (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#does-it-actually-work) 查看梯度基准测试。 Enzyme 在 LLVM IR（中间表示）级别工作，而不是在源码级别，因此任何能编译为 LLVM IR 的语言（C、C\+\+、Rust、Fortran）都可以被微分。对于 Fortran，这意味着需要选择一个前端。显而易见的选择是 Flang，LLVM 官方的 Fortran 编译器，但我们选择使用LFortran (https://lfortran.org/)，因为它对 Enzyme 来说更容易处理 (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#why-lfortran)。确定了前端后，整个链路由六个 `opt`/`clang` 调用组成，出奇地简单（尽管出问题时你最终会阅读 IR (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#what-enzyme-actually-does-to-the-ir) 来找出原因）。 编译管道：Fortran → Enzyme AD → 共享库六步编译管道，从 Fortran 源码到可微分的共享库。LFortran 将求解器降级为 LLVM IR，链接一个薄 C 包装器，Enzyme 合成导数代码，最后的优化步骤生成一个单独的 `.so` 文件，暴露正向、JVP 和 VJP 入口点。整个链在 `tesseract build` 期间运行大约 30 秒。¶ (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#id3) 六个步骤，按顺序： `` # 编译管道：Fortran 源码 → 微分后的共享库 # 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 IR 转储）、步骤 2（Enzyme 之前的 `-O1` 清理）和步骤 5（Enzyme 阶段本身），每一步都对应一个值得深思的决策。 **为什么用 LFortran 而不是 Flang。** 第一步（`lfortran --show-llvm`）也是使整个事情变得可行的原因，因为 LFortran 生成非常干净的 IR。数组以普通指针的形式出现，使用常见的 GEP/load/store 模式，很像 C 语言，而不是像 Flang（默认选项）产生的多字段描述符结构和运行时调用。这一点至关重要，因为 Enzyme 必须追踪每个内存访问以确定哪些是活跃的，而它处理指针算术比处理不透明的 `fir.box` 描述符要容易得多。 那么，是什么让 LFortran IR 变得“干净”呢？具体来说有三个因素。每个数组引用都降级为单纯的 `getelementptr` + `load`/`store` 操作，作用于裸指针；这个有 19 个参数的子程序接受 19 个普通的 `ptr` 参数，没有任何描述符结构包裹它们；整个模块包含零个函数调用：没有 `_lfortran_malloc`，没有运行时辅助函数，没有不透明的东西需要 Enzyme 追踪。下面是 LFortran 为模板体开头部分生成的 IR，涵盖了上面源码中的 `idx = (j-1)*nx + i` 索引和 `T_c = T_cur(idx)` / `kx_east = ...` 行： `` ; LFortran 为模板体生成的内容：普通指针运算，无描述符 ; 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, 乘积, 和 ... %130 = fdiv double %118, %129 ; 调和平均 store double %130, ptr %kx_east `` 就这样。数组访问是 `getelementptr` + `load`，算术是一连串的 `fmul`/`fadd`/`fdiv`，与 Fortran 逐行对应。相比之下，Flang 会将相同的数组包裹在描述符结构中，并调用运行时辅助函数，Enzyme 必须看透每一个。 缺点是 LFortran 仍在发展中，因此你的代码必须在其支持范围内 (https://lfortran.org/progress/)。 **为什么用 `-O1` 而不是 `-O3`。** 我们的第一个管道在这里使用了 `-O3`，这让我们花了数小时：正向传递完美，但 VJP 在某些输入上返回 NaN。在 `-O3` 级别，LLVM 的激进向量化和代码移动会生成 Enzyme 在反向模式下似乎会误处理的 IR 模式。在我们的案例中，当相邻单元温度相等时，中间项会相互抵消，而编译器重新排列导致除以零。解决方法是在 AD 阶段之前保持温和优化，并将 `-O3` 保留在之后。 **C 桥接。** 在 Fortran 和 Enzyme 之间有一个薄 C 包装器，它将 Fortran 的按指针 ABI 桥接到 Enzyme 可以注解的 C 接口。它本身不做任何数值工作：它在堆上分配工作数组（避免 IR 中的 `_lfortran_malloc` 调用），并通过 Enzyme 的内在函数标记哪些参数获得影子缓冲区。这个配方是通用的，不仅限于这个求解器，因此 Fortran 代码确实保持原封不动。 `` // C 桥接：在堆上分配工作数组并告诉 Enzyme 要微分什么 void thermal_2d_vjp(/* ... nx, ny, n_steps ... */) { double* T_cur = calloc(n, sizeof(double)); // 堆分配，非可分配数组 /* ... */ __enzyme_autodiff((void*)thermal_2d_solve, enzyme_const, &nx_, // 不微分 enzyme_dup, (double*)T_init, dT_init, // 值 + 影子缓冲区 enzyme_dup, &k0_, dk0, /* ... */); } `` **通过 Enzyme 进行 AD 追踪。** Enzyme 阶段（`-passes=enzyme`）是真正进行自动微分工作的地方。它分析链接后的 IR 并合成导数代码，使用存储所有磁带策略，在每个时间步缓存中间值。当它工作时，你会免费得到整个时间步进循环的伴随。当它不工作时，你将在凌晨 2 点阅读 IR 差异（参见下一步计划 (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#whats-next-and-what-wed-do-differently)）。 最后的 `-O3 + clang -shared` 步骤生成一个单独的 `.so` 文件，包含三个入口点（正向、JVP、VJP），可通过 ctypes 从 Python 调用。整个管道在 `tesseract build` 期间运行，大约需要 30 秒。 ### Enzyme 实际上对 IR 做了什么¶ (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#what-enzyme-actually-does-to-the-ir) 让我们看看 Enzyme 对 IR 做了什么。此时，干净的 LFortran IR (https://docs.pasteurlabs.ai/projects/tesseract-core/latest/blog/2026-07-09-enzyme-lfortran-autodiff.html#why-lfortran) 已经与 C 包装器链接成一个模块，这就是 Enzyme 的输入。进入 Enzyme（步骤 4 输出）时，VJP 包装器只是一个普通的调用：`__enzyme_autodiff` 内在函数，接收一个指向求解器的指针，每个参数前有一个 `enzyme_const`/`enzyme_dup` 标记（此处从全局变量加载为 `%50`/`%51`）。每个 `dup` 参数是一个值/影子对，接收其梯度的缓冲区紧随其后： `` ; combined.ll — 在 Enzyme 阶段之前 %50 = load i32, ptr @enzyme_const ; enzyme_const 标记 %51 = load i32, ptr @enzyme_dup ; enzyme_dup 标记 call void (ptr, ...) @__enzyme_autodiff( ptr nonnull @thermal_2d_solve, ; 要微分的函数 i32 %50, ptr nonnull %33, ; const: n_steps (无影子) i32 %51, ptr %3, ptr %4, ; dup: T_init, dT_init (影子) i32 %51, ptr %5, ptr %6, ; dup: k0, dk0 ; ... 另外 15 个 dup 对 ... i32 %51, ptr nonnull %43, ptr %28) `` 经过阶段（步骤 5 输出）后，这个内在函数消失了，`__enzyme_autodiff` 在整个模块中除了现在未使用的声明外无处可见。取而代之的是 Enzyme 从头合成的函数 `@diffethermal_2d_solve`。它的签名是线索：原始求解器中的每个活跃参数现在都出现两次，值后面紧跟一个带有 `'` 后缀的影子，它将接收其梯度。 `` ; ad.ll — 合成函数的签名 define internal v