NVIDIA Warp加速微分物理仿真技术

原创

用户11764306

发布于 2026-06-07 23:05:34

使用NVIDIA Warp构建用于AI的加速、可微分计算物理代码

2026年3月12日

作者：Sheel Nidhan, Eric Shi, Neil Ashton, Zach Corse 和 Mohammad Mohajerani

引言

计算机辅助工程正在从人工驱动的工作流转向 AI 驱动的工作流，包括能够在不同几何形状和运行条件下泛化的物理基础模型。与大型语言模型不同，这些模型依赖于大量高保真、符合物理规律的数据。

最近关于计算流体动力学代理模型的规模定律研究表明，仿真生成的训练数据在实践中往往是成本的主要限制因素。这给模拟器提出了更高的要求：必须原生支持 GPU、运行快速，并且能够直接接入机器学习工作流。

NVIDIA Warp 是一个用于加速仿真、数据生成和空间计算的框架，它架起了 CUDA 和 Python 之间的桥梁。Warp 使开发者能够将高性能内核编写为普通的 Python 函数，这些函数会被即时编译为在 GPU 上执行的高效代码。与基于张量的框架（开发者需要将计算表达为对整个 N 维数组的操作）不同，在 Warp 框架中，开发者编写灵活的内核，这些内核会在计算网格的所有元素上同时执行。

仿真内核通常在计算网格上表达，并依赖于数据相关的控制流，例如条件分支、提前退出以及按元素变化的更新。在张量框架中，这些模式需要组合布尔掩码，这不仅很快就会变得难以处理，还可能在与无关元素上浪费计算。在 Warp 内核中，每个线程可以独立分支、跳过或退出，自然表达这些逻辑，而无需掩码变通方案。

此外，本文将展示，借助 Warp 原生支持的自动微分，用 Warp 编写的求解器可以轻松实现可微分，并且能够直接集成到优化或训练工作流中，同时保持与 PyTorch、JAX 和 NumPy 等框架的互操作性，适用于仿真、机器人、感知和几何处理等应用场景。

本文将指导你如何使用 Warp 从头构建一个二维纳维-斯托克斯求解器，解释 Warp 编程模型如何映射到偏微分方程求解器，然后通过对仿真进行微分来解决一个端到端的最优扰动问题。最后通过工业案例研究展示 Warp 在生产工作流中的能力。更多信息请参见 NVIDIA/warp GitHub 仓库中的二维纳维-斯托克斯求解器示例和二维纳维-斯托克斯最优扰动示例。

如何使用 Warp 编写二维纳维-斯托克斯求解器

为了将重点放在 Warp 而非数值方法上，这里采用一个教科书式的二维衰减湍流示例，由不可压缩纳维-斯托克斯方程的涡量-流函数形式描述。涡量 𝜔 根据输运方程演化：

∂𝜔/∂𝑡 + (∂𝜓/∂𝑦)·(∂𝜔/∂𝑥) – (∂𝜓/∂𝑥)·(∂𝜔/∂𝑦) = (1/Re) ∇²𝜔

流函数 𝜓 通过泊松方程从涡量恢复：

∇²𝜓 = –𝜔

在周期性边界条件下，上述方程在傅里叶空间中简化为代数方程，无需迭代求解器：

ˆ𝜓{m,n} = ˆ𝜔{m,n} / (𝑘²𝑥 + 𝑘²𝑦)

其中 (𝑘𝑥, 𝑘𝑦) 是傅里叶空间中的波数对。求解器利用快速傅里叶变换算法高效地将 𝜔 和 𝜓 变换到傅里叶空间再变换回来。

每个时间步包含两个子步骤（图1）。首先，在 𝐿×𝐿 方形区域内的 𝑁×𝑁 网格上对涡量输运方程进行离散化，使用三阶强稳定性保持龙格-库塔（RK3）方案将解向前推进 Δ𝑡，得到 𝜔(𝑡+Δ𝑡)。其次，在傅里叶空间中求解泊松方程，得到更新后的 𝜓(𝑡+Δ𝑡)。

图1. 求解器单个时间步循环示意图

构建块1：有限差分离散化与时间推进

涡量输运方程中的对流项和扩散项采用二阶中心有限差分近似。高阶离散化也可使用，但为简单起见选择中心二阶格式。

图2. 𝜔 和 𝜓 的有限差分模板

以下 rk3_update() 内核计算扩散项和对流项，并执行单个 RK3 子步更新。step() 函数在每个时间步调用该内核三次，每个 RK3 阶段一次，每个阶段使用不同的系数。

@wp.kernel
def rk3_update(
    n: int, h: float, re: float, dt: float,
    coeff0: float, coeff1: float, coeff2: float,
    omega_0: wp.array2d(dtype=float),
    omega_1: wp.array2d(dtype=float),
    psi: wp.array2d(dtype=float),
    omega_out: wp.array2d(dtype=float)):
    
    i, j = wp.tid()
    left = cyclic_index(i - 1, n)
    right = cyclic_index(i + 1, n)
    top = cyclic_index(j + 1, n)
    down = cyclic_index(j - 1, n)
    
    inv_h2 = 1.0 / (h * h)
    laplacian = (omega_1[right, j] + omega_1[left, j] + 
                 omega_1[i, top] + omega_1[i, down] - 4.0 * omega_1[i,j]) * inv_h2
    
    inv_2h = 1.0 / (2.0 * h)
    j1 = ((omega_1[right, j] - omega_1[left, j]) * inv_2h) * \
         ((psi[i, top] - psi[i, down]) * inv_2h)
    j2 = ((omega_1[i, top] - omega_1[i, down]) * inv_2h) * \
         ((psi[right, j] - psi[left, j]) * inv_2h)
    
    rhs = (1.0 / re) * laplacian + j2 - j1
    omega_out[i, j] = coeff0 * omega_0[i, j] + coeff1 * omega_1[i, j] + coeff2 * dt * rhs

rk3_update() 内核遵循单指令多线程范式，其中每个线程映射到计算域上的一个网格点，所有 𝑁×𝑁 个点通过一次 wp.launch() 调用同时更新。

wp.launch(rk3_update,
          dim=(self.n, self.n),
          inputs=[self.n, self.h, self.re, self.dt,
                  stage_coeff[0], stage_coeff[1], stage_coeff[2],
                  self.omega_0, 
                  self.omega_1, 
                  self.psi],
          outputs=[self.omega_tmp])

图3. 二维网格上 𝜔 的 SIMT 更新。线程 (i, j) 使用当前时间步模板中相邻单元格的值将单元格 (i, j) 更新到下一个时间步

构建块2：FFT 泊松求解器

Warp 基于图块的原语支持在傅里叶空间中求解泊松方程。关键操作是 wp.tile_fft() 和 wp.tile_ifft()，它们分别对加载到图块中的单行执行正向和反向 FFT。一个 𝑁×𝑁 数组上的完整二维 FFT 分解为三个步骤：逐行 FFT → 转置 → 逐行 FFT。图4 说明了 fft_tiled() 和 ifft_tiled() 如何计算正向和反向 FFT。

@wp.kernel
def fft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)):
    i, _, _ = wp.tid()
    a = wp.tile_load(x, shape=(1, N_GRID), offset=(i, 0))
    wp.tile_fft(a)
    wp.tile_store(y, a, offset=(i, 0))

@wp.kernel
def ifft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)):
    i, _, _ = wp.tid()
    a = wp.tile_load(x, shape=(1, N_GRID), offset=(i, 0))
    wp.tile_ifft(a)
    wp.tile_store(y, a, offset=(i, 0))

图4. 在 NxN 网格上逐行 tile_fft。每个块将一行加载到寄存器图块中，协作计算 FFT，然后将结果存回全局内存

二维 FFT 还需要在逐行变换之间进行一次转置。可以使用 SIMT 或图块范式。为简单起见，下面展示 SIMT 版本：

@wp.kernel
def transpose(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)):
    i, j = wp.tid()
    y[i, j] = x[j, i]

将这三个内核组合起来：fft_tiled -> transpose -> fft_tiled，即可得到完整的二维正向 FFT。逆向变换使用 ifft_tiled 遵循相同模式。

整合构建块

示例中的 step() 函数依赖一些其他辅助内核。有了所有构建块，一次 step() 调用将仿真推进一个时间步。示例代码中的 self._solve_poisson() 方法将 𝜔(𝑡+Δ𝑡) → FFT → ˆ𝜔 → (方程3) → ˆ𝜓 → IFFT → 𝜓(𝑡+Δ𝑡) 流程封装为模块化调用。

def step(self) -> None:
    for stage_coeff in self.rk3_coeffs:
        wp.launch(
            rk3_update,
            dim=(self.n, self.n),
            inputs=[self.n, self.h, self.re, self.dt,
                    stage_coeff[0], stage_coeff[1], stage_coeff[2],
                    self.omega_0,
                    self.omega_1,
                    self.psi],
            outputs=[self.omega_tmp],
        )
        self.omega_1, self.omega_tmp = self.omega_tmp, self.omega_1
        self._solve_poisson()
    wp.copy(self.omega_0, self.omega_1)

运行求解器产生图5所示的衰减湍流场。在 GPU 上，step() 函数通过 wp.ScopedCapture 被捕获到 CUDA Graph 中，后续所有帧使用 wp.capture_launch() 重放，消除每次启动的开销。

图5. Re = 1,000 时的二维衰减湍流

对求解器进行微分

现在已经构建了可工作的求解器，下一个问题是如何使其可微分。

自动微分通过对计算图中每个基本运算应用链式法则来计算程序的精确导数。与有限差分不同，AD 避免了步长调整，并产生达到机器精度的梯度。AD 对 PDE 求解器的主要优势在于可扩展性：对于大规模网格上的复杂仿真，单次正向求解已经非常昂贵，而有限差分等方法需要 𝑂(𝑛) 次完整求解才能获得相对于 𝑛 个输入的梯度。

反向模式 AD 通过大约一次正向传递加一次反向传递计算出所有 ∂𝐿/∂𝑥_𝑖，使得基于梯度的优化在生产分辨率下变得可行。这与神经网络中的反向传播思想相同，也是深度学习和大规模物理优化能够处理数百万自由度的原因。

Warp 自动微分系统在编译时为一个可微分仿真生成两个版本的程序：

正向版本：接收物理输入（初始条件、离散化控制律等）并计算仿真输出（场量、导出量）以及伴随版本所需的中间数组。
伴随版本：正向仿真的自动生成对应部分，接收所选目标量相对于仿真输出的敏感性，并将其一直反向传播到输入。这种反向传播复用正向执行中的中间数组，在整个求解器上应用微分链式法则，无需构建大的符号表达式即可产生仿真伴随。

开发者编写正向物理过程，Warp 负责梯度计算。任何需要可微分的 wp.array 在分配时设置 requires_grad=True，这告诉 Warp 分配一个伴随存储数组。生成的伴随可以独立使用，或与 PyTorch、JAX 互操作以实现端到端优化，包括训练机器学习模型。目前，Warp 仅支持反向模式 AD。

为了说明，本文解决了某文献中提出的最优扰动问题。在湍流中，初始条件的微小扰动会随时间放大，并显著改变流动的轨迹。识别哪些扰动增长最快是迈向流动控制和理解哪些流动结构具有动态重要性的基石。具体来说，目标是寻找初始涡量扰动 Δ𝜔，使得在提前时间 𝜏 内扰动轨迹与未扰动轨迹之间的差异最大化。

设 𝐹𝜏 表示应用于 𝜏 时间单位的前向求解器。未扰动轨迹为 𝑌∗ = 𝐹𝜏(𝜔₀)，扰动轨迹为 ˜𝑌 = 𝐹_𝜏(𝜔₀ + Δ𝜔)。最小化均方误差：

MSE = –(1/𝑁²)·‖𝑌∗ – ˜𝑌‖₂²

其中负号将最大化轨迹差异转化为最小化问题。为了约束优化，要求 Δ𝜔 的均方根 ≤ 初始涡量场 𝜔₀ 均方根的 20%。

无就地修改

wp.Tape() 记录前向传递中的内核启动，并在反向传递中重放它们以计算梯度。这仅在反向传递所需的中间值仍然可用时才有效，因此数组不能随意就地覆盖。这是与不可微分求解器的关键区别。在仅正向版本中，可以在每个时间步结束时交换两个数组 omega_0 和 omega_1。对于可微分求解器，需要将 RHS 计算和 RK3 更新拆分为写入不同数组的独立内核。

在 Warp 中，所有中间数组都需要由用户显式定义。这需要为每个时间步的每个 RK 子步预先分配单独的数组，这是任何可微分求解器的主要 GPU 内存成本。

self.omega_timestep = [wp.zeros((n, n), dtype=wp.float32, requires_grad=True) for _ in range(T + 1)]

self.omega_stage = []
self.psi_stage = []
self.rhs_stage = []
self.fft_arrays = []

for _ in range(T):
    s_omega, s_psi, s_rhs, s_fft = [], [], [], []
    for _ in range(3):
        s_omega.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True))
        s_psi.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True))
        s_rhs.append(wp.zeros((n, n), dtype=wp.float32, requires_grad=True))
        s_fft.append({"omega_complex": wp.zeros((n, n), dtype=wp.vec2f, requires_grad=True)})
    self.omega_stage.append(s_omega)
    self.psi_stage.append(s_psi)
    self.rhs_stage.append(s_rhs)
    self.fft_arrays.append(s_fft)

为每个中间状态存储 Warp 数组会随时间步数线性增长，在长时间运行中变得不可行。一种常见方法是梯度检查点：仅保存选定的状态，然后在反向传递期间使用正向求解器重新计算缺失的片段。这种方法用额外的正向计算换取更小的内存占用。

使用 wp.Tape() 记录梯度

有了预分配的数组，记录和微分正向传递非常简单：

with wp.Tape() as tape:
    forward()

tape.backward(loss)

wp.Tape() 上下文将每个 wp.launch() 调用记录到一个计算图中。tape.backward(loss) 反向遍历该图，计算损失相对于 Warp 数组的导数。

优化循环

以下代码展示了一个优化步骤。forward() 函数在扰动的初始涡量上运行，生成最终场和损失。优化器更新扰动以减少损失。

with wp.Tape() as tape:
    forward()

tape.backward(loss)
optimizer.step([delta_omega.grad.flatten()])
tape.zero()

经过 1000 次迭代后，优化器发现了一个结构化的扰动 Δ𝜔，使轨迹差异放大，MSE 从接近零上升到约 250。

图6. 1000 次迭代中的优化过程及发现的扰动（右上）

Warp 实践：AI 驱动工业工作流的案例研究

某机构 XLB

某机构构建了 XLB，一个具有 Warp 和 JAX 两种后端的可微分格子玻尔兹曼求解器。在约 1.34 亿单元的顶盖驱动空腔基准测试中，Warp 在单个 40 GB 某 GPU 上比 JAX 快约 8 倍，大致相当于 JAX 需要 8 个某 GPU 才能达到的吞吐量。在更大规模上，Warp 内存使用减少约 2.5 到 3 倍，并完成了 JAX 在同一 GPU 上内存不足的最大案例。

图7. Warp 与 JAX 的吞吐量和内存使用对比