计算机辅助工程 (CAE) 正在从人工驱动的工作流转向 AI 驱动的工作流，包括跨几何图形和操作条件进行泛化的物理基础模型。与 LLM 不同，这些模型依赖于大量高保真、物理合规的数据。
近期关于计算流体动力学 (CFD) 代理的扩展定律研究表明，仿真生成的训练数据在实践中通常会限制成本。这会将要求推送到模拟器上，模拟器必须原生、快速且能够直接插入 ML 工作流。
NVIDIA Warp 是一个用于加速仿真、数据生成和空间计算的框架，将 CUDA 和 Python 连接在一起。Warp 使开发者能够将高性能核函数编写为常规 Python 函数，并将其 JIT 编译为可在 GPU 上执行的高效代码。与基于张量的框架 (开发者将计算表示为对整个 N 维数组的运算) 不同，开发者在 Warp 框架中编写了灵活的核函数，可在计算网格的所有元素上同时执行。
模拟内核通常在计算网格上表示，并且依赖于数据相关的控制流，例如条件、早期退出和选择性更新 (每个元素各不相同) 。在张量框架中，这些模式需要合成布尔掩码，而这很快就变得难以操作，并可能会在不相关的元素上浪费计算。在 Warp 内核中，每个线程都可以独立分支、跳过或退出，从而自然地表达此逻辑，而无需屏蔽变通方案。
此外，正如本文所展示的，使用 Warp 编写的求解器可以通过对自动微分的 Warp 原生支持轻松微化。它们易于与优化或训练工作流集成，同时还能与 PyTorch、JAX 和 NumPy 等框架进行互操作，适用于仿真、机器人、感知和几何处理等用例。
本文将为您介绍如何完全在 Warp 中构建 2D Navier – Stokes 求解器。它解释了 Warp 编程模型如何映射到 PDE 求解器上。然后，它通过模拟进行区分，以端到端解决最佳微扰问题。最后，行业案例研究展示了 Warp 在生产工作流中可以实现的功能。有关更多信息，请在NVIDIA/warp GitHub 库上参阅 2D Navier–Stokes 求解器示例 和 2D Navier-Stokes 最优微扰示例。
如何编写 2D Navier–使用 Warp 的 Stokes 求解器
为了始终专注于 Warp 而不是数值方法，此处使用了 2D 衰减流的教科书示例，该示例通过不可压缩纳维尔 – 斯托克斯方程的涡流函数表述进行了描述。涡量 \(\omega\) 根据传输方程演变：
\(\frac{\partial \omega}{\partial t} + \frac{\partial \psi}{\partial y}\frac{\partial \omega}{\partial x} – \frac{\partial \psi}{\partial x}\frac{\partial \omega}{\partial y} = \frac{1}{\text{Re}}\nabla^2 \omega \tag{1}\)
流函数 \(\psi\) 通过泊松方程从涡流中恢复：
\(\nabla^2 \psi = -\omega \tag{2}\)
在周期边界条件下，上述方程简化为里叶空间中的代数方程，而无需迭代求解器：
\(\hat{\psi}_{m,n} = \frac{\hat{\omega}_{m,n}}{k_x^2 + k_y^2} \tag{3}\)
其中 \((k_x, k_y)\) 是里叶空间中的波数对。求解器利用快速里叶变换 (FFT) 算法将 \(\omega\) 和 \(\psi\) 高效转换为里叶空间，反之亦然。
每个时间步长有两个子组件 (图 1) 。首先，涡量传输方程在 \(L \times L\) 平方域上的 \(N \times N\) 网格上离散。使用保持三阶强稳定性的 Runge-Kutta (RK3) 方案，该解决方案会在 \(\Delta t\) 之前及时向前推进，以获得 \(\omega(t+\Delta t)\)。其次，在里叶空间中求解泊松方程，以获得更新后的 \(\psi(t+\Delta t)\)。
因此，前向求解器具有两个构建块，后续章节中将对此进行介绍：
- 用于离散化和时间行进的 Warp 内核
- 基于 FFT 的泊松求解器
基础模组 1：有限差分离散和时间行进
图 2 显示了涡量传输方程中的对流和扩散项与二阶中心有限差的近似值。也可以使用高阶离散，但为简单起见，我们选择了中心二阶方案。
以下内容rk3_update ()内核计算扩散和对流项，并执行单个 RK3 子步骤更新。这个step ()函数在每个时间步调用此核函数三次，在每个 RK3 阶段调用一次，其系数不同 (coeff0、coeff1、coeff2) 每个阶段。
@wp.kerneldef 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)): """Perform a single substep of SSP-RK3.""" 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 () 核函数遵循单指令多线程 ( SIMT) 范式，即每个线程映射到计算域上的一个网格点，并且所有 \(N \times N\) 点都通过一次 wp.launch () 调用同时更新。
wp.launch(rk3_update, dim=(self.n, self.n), # one thread per grid point 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] ) |
构建块 2：FFT 泊松求解器
基于 Warp Tile 的基元能够在里叶空间中求解泊松方程。关键运算为 wp.tile_fft () 和 wp.tile_ifft ()，它们分别在加载到图块中的单行上执行正向和反向 FFT。然后，将 \(N \times N\) 阵列上的完整 2D FFT 分解为三个步骤：逐行 FFT – > 转置 – > 逐行 FFT。图 4 中的原理图解释了 fft_tiled () 和 ifft_tiled () 如何在底层计算正向和反向 FFT。
@wp.kerneldef fft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)): """Row-wise FFT using tile primitives.""" 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.kerneldef ifft_tiled(x: wp.array2d(dtype=wp.vec2f), y: wp.array2d(dtype=wp.vec2f)): """Row-wise inverse FFT using tile primitives.""" 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)) |
tile_fftNxN 网格。每个块将一行加载到寄存器图块中，协同计算 FFT，并将结果存储回全局内存2D FFT 还需要在行级通道之间进行转置。这可以使用 SIMT 或 tile 范例 (通过wp.tile_transpose) 。为简单起见，SIMT 版本如下所示：
@wp.kerneldef 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– >转置– >fft_tiled它们共同构成了完整的 2D 前向 FFT。与之相反的模式与ifft_tiled。
将基础模组组合在一起
此示例中的 step() 函数依赖于此处未详细讨论的其他一些辅助内核。有关这些内核的定义，请参阅 NVIDIA/warp GitHub 库上的 2D Navier–Stokes 求解器示例。完成所有构建块后，只需执行一个 step() 调用，即可将模拟推进一个时间步。示例代码中的 self._solve_poisson() 方法将 \(\omega(t+\Delta t) \xrightarrow{\text{FFT}} \hat{\omega} \xrightarrow{\text{Eq.\,3}} \hat{\psi} \xrightarrow{\text{IFFT}} \psi(t+\Delta t)\) 工作流抽象化，实现模块化。
def step(self) -> None: """Advance simulation by one timestep using SSP-RK3.""" 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], ) # Swap buffers for next RK3 substep self.omega_1, self.omega_tmp = self.omega_tmp, self.omega_1 # Update streamfunction for next timestep self._solve_poisson() # Copy updated vorticity to self.omega_0 for the next timestep wp.copy(self.omega_0, self.omega_1) |
运行求解器会产生如图 5 所示的衰减流场。在 GPU 上，step() 函数通过 wp.ScopedCapture 捕获到 CUDA Graph 中，并使用 wp.capture_launch() 对所有后续帧进行回放，从而消除每次启动的开销。
通过求解器进行区分
现在，工作求解器已经构建完成，下一个问题是如何使其可微分。
自动微分(AD) 通过对计算图形中的每个基本运算应用链式法则来计算程序的精确导数。与有限差分不同，AD 可避免步长调整，并生成符合机器精度的梯度。AD 对于 PDE 求解器的主要优势在于扩展：在大型网格上进行复杂模拟时，每个前向求解都成本高昂，因此有限差值等方法需要 \(O(n)\) 全求解才能获得 \(n\) 输入的梯度。
反向模式 AD 大约以一次正向传递和一次反向传递的方式计算所有 \(\partial \mathcal{L}/\partial x_i\)，这使得基于梯度的优化在生产分辨率方面切实可行。这与神经网络中的反向传播理念相同，这就是深度学习和大规模物理优化可以处理数百万自由度的原因。
Warp 自动微分系统在编译时为可微分模拟生成两个版本的程序：
- 前向版本: 获取物理输入 (初始条件、离散治理定律等) 并计算模拟输出 (字段、衍生数量) 以及伴随版本所需的中间数组的代码。
- 伴随版本: 与前向模拟自动生成的对应版本，可以获取所选感兴趣数量相对于模拟输出的敏感度，并将其传播回输入。这种反向传播会重用前向执行中的中间数组，在整个求解器中应用微分链式规则，从而在不构建大型符号表达式的情况下生成模拟伴随。
开发者编写正向物理，Warp 负责处理梯度计算。任何wp.array应该是可微分的requires_grad = True该命令指示 Warp 为伴随存储分配一个配套数组。生成的伴随可以独立使用 (如本示例所示) ，也可以与 PyTorch 或 JAX 进行互操作，以实现端到端优化，包括训练 ML 模型。目前，Warp 仅支持反向模式 AD。
为说明这一点，我们将在此处讨论使用生成式对抗网络预测和控制二维衰减流中概述的最佳微扰问题。在流中，对初始条件的微小干扰会随着时间的推移而放大，并显著改变流动的轨迹。识别哪些微扰增长最快，是实现流量控制和了解流量中哪些结构具有动态显著性的基础。具体而言，我们要寻找初始涡量微扰 \(\Delta\omega\)，该微扰可在提前时间 \(\tau\) 最大限度地增大受干扰和未受干扰轨迹之间的差异。
\(F^{\tau}\) 表示应用于 \(\tau\) 时间单位的前向求解器。未受干扰的轨迹为 \(Y^{*} = F^{\tau}(\omega_0)\)，受干扰的轨迹为 \(\tilde{Y} = F^{\tau}(\omega_0 + \Delta\omega)\)。均方误差 (MSE)
\(\mathrm{MSE} = -\frac{1}{N^2}\left\| Y^* – \tilde{Y} \right\|_2^2 \tag{4}\)
最小化，其中负号将轨迹差异的最大化转化为最小化问题。为了约束优化 \(\mathrm{rms}(\Delta\omega) \leq 0.2 \times \mathrm{rms}(\omega_0)\)，即干扰 RMS 不得超过初始涡量场 \(\omega_0\) 的 RMS 的 20%。
有关更多详细信息，请参阅 2D Navier-Stokes 最佳微扰示例 on the NVIDIA/warp GitHub 库。以下章节将重点介绍前向求解器中使其可微分的三项关键更改。
无需就地修改
wp.Tape ()记录核函数在正向传递中的启动，并反向回放以计算梯度。只有在反向传递所需的中间值仍然可用时，这种方法才有效，因此数组无法在原位自由覆盖。这是与不可微分求解器的关键区别。在前向版本中，可以切换两个数组，omega_0以及omega_1，在每个时间步结束时：
wp.copy(omega_0, omega_1) |
对于可微求解器，RHS 计算和 RK3 更新需要拆分成单独的内核，这些内核会写入单独的数组。因此，单个 RK3 更新将变为以下内容。请注意omega_1值无法复制到omega_0和之前一样在每个时间步的末尾。
omega_out[i, j] = coeff0 * omega_0[i, j] + coeff1 * omega_in[i, j] + coeff2 * dt * rhs[i, j] |
在 Warp 中，用户需要显式定义所有中间数组。这需要在每个时间步为每个 RK 子步骤预先分配单独的数组，这通常是任何可微分求解器的主要 GPU 显存成本。
self.omega_timestep = [wp.zeros((n, n), dtype=wp.float32, requires_grad=True) for _ in range(T + 1)]# Intermediate arrays for each RK3 substep for each timestepself.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), # ... plus 4 FFT scratch arrays, each (n, n) vec2f }) 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 数组会随时间步长的数量呈线性增长，这在长期运行中变得令人望而却步。一种常见方法是梯度检查点，仅保存选定的状态，然后在向后传递期间使用前向求解器重新计算缺失的段。这种方法以更小的内存占用来交换额外的前向计算。有关如何在 Warp 中实现梯度检查点的示例，请参阅 流体检查点示例 on the NVIDIA/warp GitHub 存储库。
使用wp.Tape()
在预先分配的数组就位后，记录和区分前向传递非常简单：
with wp.Tape() as tape: forward() # wp.launch calls that take omega from t0 to t0 + lead t and calculate MSE tape.backward(loss) # Automatic differentiation to get derivatives of loss w.r.t Warp arrays |
这个wp.Tape ()上下文记录wp.launch ()对计算图形进行调用。tape.backward (损失)反向遍历该图形，计算损失Warp 数组。这里的重点是损失对于 \(\Delta{\omega}\)，可以通过delta_omega.grad。
优化循环
以下代码块显示了一个优化步骤。这个forward ()该函数在受干扰的初始涡上运行，以产生最终场和损失 ( MSE 与未受干扰的运行) 。磁带记录核函数在此过程中的启动。tape.backward (损失)然后通过记录的图形进行反向传播，以计算与微扰相关的梯度optimizer.step ()更新微扰以减少损失。最后，tape.zero ()在下一次迭代之前清除累积的梯度。
with wp.Tape() as tape: forward() # Loss is computed inside forward() functiontape.backward(loss)optimizer.step([delta_omega.grad.flatten()])tape.zero() |
经过 1000 次迭代后，优化器发现了结构化微扰 \(\Delta\omega\)，该微扰放大了轨迹离散，将 MSE 从近零推升至约 250。从求解器在环优化中获得的微扰场在质量上类似于使用生成式对抗网络预测和控制二维衰减流中报告的干扰场，详见Prediction and Control of Two-Dimensional Decaying Turbulence Using Generative Adversarial Networks。
![Optimization GIF for 1,000 iterations. $\mathrm{MSE}(Y^*,\tilde{Y}$ decreases over iterations (left), with field snapshots showing the baseline $\omega_0$ (top center), learned perturbation $\Delta\omega$ (top right), target $Y^*$ (bottom center), and optimized output $\tilde{Y}$ (bottom right) on a $[0,2\pi]\times[0,2\pi]$ domain.](https://developer.nvidia.cn/blog/wp-content/uploads/2026/02/optimization-progressing-1000-iterations-discovered-perturbation-1.gif)
如需了解更多信息，NVIDIA/ Warp GitHub 资源库包含 CFD 以外的其他可微分求解器示例。另请参阅不断增长的 利用 Warp 的研究出版物列表。
Warp 实际应用：AI 驱动的工业工作流案例研究
在真实的 AI 工作流中，仿真和几何图形位于更大的系统 (代理模型、RL、设计优化等) 中。PyTorch 和 JAX 可处理训练和张量运算，但仿真层添加了分段时间步、模板更新和大型空间查询。Warp 的目标是内核密集型层：您可以控制执行、融合内核以减少内存流量并启动，以及使用 CUDA 计算图减少重复调度。它还能与 PyTorch 和 JAX 张量进行零复制互操作。
Autodesk XLB
Autodesk Research 构建了 XLB，这是 Python 中的一个可微分的格子玻尔兹曼求解器，具有 Warp 和 JAX 后端，可以直接比较相同的公式和硬件。在价值 1.34 亿单元的盖子驱动型腔基准测试中，Warp 在单个 40 GB NVIDIA A100 Tensor Core GPU 上的运行速度比 JAX 快 8 倍，大致相当于 JAX 需要 8 个 A100 Tensor Core GPU 才能达到的吞吐量。在更大的尺寸下，Warp 使用的内存大约减少了 2.5 倍到 3 倍，并且完成了最大的情况，在同一 GPU 上 JAX 显存耗尽。
如需了解详情，请参阅 Autodesk Research 在 NVIDIA GH200 上为计算流体动力学带来扭曲速度。
Google DeepMind MuJoCo
Google DeepMind 最近发布了基于 Warp 的大规模多体动力学后端 MuJoCo Warp (MJWarp) 。在同等硬件上，与 JAX 相比，Warp 后端实现了高达 252 倍 (运动) 和 475 倍 (操作) 的加速。MJWarp 通过利用稀疏矩阵运算和推理执行更精确地调度计算来实现这一目标，同时保持与 JAX 训练兼容的插件。
如需了解详情，请参阅 MuJoCo Warp 发布公告。
C-Infinity AutoAssembler
C-Infinity AutoAssembler ASI 引擎展示了 Warp 在物理模拟之外的 AI 驱动工业工作流中的价值。它通过直接从原始几何图形中计算接触、干扰和间隙，将全保真 CAD 组件转换为用于 AI 规划的运动约束。当前的 CAD 系统不支持这些关键查询，因为构建制造流程计划、评估设计变更和生成执行指令都需要这些查询。
AutoAssembler ASI 引擎支持构建制造编译器，将工程 CAD 数据直接转换为供人类或机器人使用的装配指令。该技术使用针对大规模处理优化的 Warp 内核实施，以构建空间智能。
与优化的 CPU 基准相比，在 NVIDIA L4 Tensor Core GPU 上，Warp GPU 后端实现了高达 669 倍的加速 (基于 FCL 和 Embree 等先进库) 。该技术已经在顶级 OEM 的企业制造工作流中使用。
如需了解详情，请参阅 AutoAssembler ASI：加速空间智能、C-Infinity。
开始使用适用于计算物理应用的 Warp
借助 Warp，您可以在 Python 中将物理和几何体编写为 GPU 内核，而无需将所有内容强制引入基于张量的框架。在 CFD 中，时间步长和可微分可将映射清晰地解析为内核，保持物理结构的完整性。
该模型已经出现在工业 AI 工作流中，包括 Autodesk 可微 CFD 求解器、Google DeepMind 多体动力学工作和 C-Infinity 空间推理引擎。通过与 PyTorch 和 JAX 的零拷贝互操作，Warp 可插入 ML 流程，同时保留这些工作负载所需的控制流，并在性能、内存和可扩展性方面实现显著提升。
要开始将 Warp 用于计算物理应用，请查看以下资源：
如需了解更多信息，请参加 NVIDIA GTC 2026 会议：如何使用 NVIDIA Warp 构建 GPU 加速的计算物理模拟【DLIT81837】。观看 GTC 主题演讲 与 NVIDIA 创始人兼首席执行官黄仁勋，探索更多 物理 AI、机器人 和 视觉 AI GTC 会议。
致谢
感谢 Felix Meyer 为此博文和项目做出的贡献。
