← Tesseract 博客 如果你能够通过现有的Fortran、C或C++模拟代码进行反向传播，将其嵌入到JAX和torch中，并将其用作高性能的可微分物理引擎，那会怎么样？事实证明，你可以——如果你足够勇敢的话……数十年的CFD、气候、航空航天和核能的经过验证的物理代码在现代机器学习管道无法跨越的墙壁后面，因为它们没有暴露梯度。通常的解决方案是将所有代码重写为JAX或PyTorch。我们在这里探索的替代方案是将代码留在原地，并且通过一些LLVM级别的魔法无论如何得到准确的梯度。这是可能的，因为Enzyme在LLVM IR级别应用自动微分，因此我们可以对任何编译为LLVM的代码进行微分！我们需要做的就是用胶带将LFortran、LLVM和Enzyme粘在一起，将结果指向一个Fortran热解算器，并从另一端获得准确的梯度。从那里，Tesseract将结果包装为自定义的JAX原语，因此Fortran解算器在任意JAX代码中成为一个可微分的层。这一切都是相当实验性的，因此要准备花一些时间追踪返回NaN的梯度，并手动比较LLVM IR差异以使其正常工作。但是如果你付出劳动，你将获得整个多步骤时间循环的梯度，这些梯度与解析解匹配，看到一个结合一些最古老和最新技术的堆栈能够协同工作是非常惊人的。但让我们从头开始。接下来是完整的故事，包括编译管道、我们遇到的尖锐边缘以及在一个具有挑战性的反问题上的应用（如果没有AD，根本不可能做到的）。从遗留代码中提取梯度的问题 ¶ 如果你在科学计算领域工作，可能会遇到这个问题：你有一个用Fortran（或C，或C++）编写的模拟，现在有人需要梯度（模拟输出相对于输入的导数）。也许是为了优化，也许是反问题，也许是将模拟嵌入到机器学习管道中。现在你的选择是：手写伴随。归结为手动实现每行代码的导数，这意味着需要几个月的专家努力。容易出错，并且维护噩梦，随着前向代码逐渐失去同步。有限差分。扰动每个输入并计算输出的差异。慢（对于n个参数需要O(n)评估），不准确，并且在刚性问题上条件较差。用JAX或PyTorch重写并使用自动微分。当然，如果你想将数万行经过验证的物理重新编写。但如果你可以自动从现有源代码编译导数呢？这就是你尝试时实际的样子。 一个Fortran热解算器的例子 ¶ 我们的测试对象是 thermal_2d.f90 ，大约220行的原始Fortran 90代码，是我们为这个实验编写的。它解决了具有温度依赖导热性的二维瞬态热传导问题： \[ ho \, c_p \frac{\partial T}{\partial t} = \nabla \cdot \big( k(T) \, \nabla T \big) + Q\] 导热系数 k 遵循线性材料模型，\(k(T) = k_0 + k_1 \, T\)，时间积分采用显式Euler法进行，步数为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) 是非线性的，模板系数依赖于当前温度场，因此雅可比在每个时间步上都会变化。手动编码伴随程序，并在每次需要时重新推导它