十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

NVIDIA Warp 上手与深入:用 Python 编写 GPU 加速的可微分模拟内核

NVIDIA Warp 上手与深入:用 Python 编写 GPU 加速的可微分模拟内核 NVIDIA Warp 上手与深入用 Python 编写 GPU 加速的可微分模拟内核【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warpWarp 是一个面向 GPU 加速模拟、机器人与机器学习的 Python 框架它把普通 Python 函数通过 JIT 即时编译成可在 CPU 或 GPU 上运行的高效内核代码并提供物理模拟、几何处理等丰富原语。读完本文你将掌握 Warp 的安装方式、平台支持范围、内核编写与启动的核心 APIwp.kernel/wp.launch/wp.array并能直接运行仓库自带的仿真示例作为自己的起点。Warp 是什么从 Python 函数到内核代码Warp 的核心思想可以用一句话概括取普通 Python 函数JIT 编译成可在 CPU 或 GPU 上运行的高效内核代码。开发者用 Python 编写逻辑Warp 在运行时将其编译为原生内核从而兼顾开发效率与计算性能。它适用于物理模拟、机器人学、几何处理等领域并且内核是可微分的——这意味着内核可以作为 PyTorch、JAX、Paddle 等机器学习框架流水线的一部分参与反向传播训练。从源码结构看这套机制的实现集中在 warp/_src/context.py 中kernel()装饰器第 1816 行起负责把带类型注解的 Python 函数注册为 Warp 内核launch()第 11266 行起负责在目标设备上异步启动内核。Warp 内置函数如wp.tid()、wp.length_sq()、wp.normalize()在 warp/_src/lang.py 中有统一说明内置函数带有Kernel可在内核内调用、Python可在 Python 作用域调用、Differentiable反向自动微分时传播梯度三类标签这也是后续阅读 Warp 文档时的重要线索。快速上手20 行代码模拟一百万粒子README 提供了一个极具代表性的 Quick Start——用约 20 行代码模拟一百万个粒子在引力作用下的运动。完整的可运行代码如下import warp as wp import numpy as np num_particles 1_000_000 dt 0.01 wp.kernel def gravity_step(pos: wp.array[wp.vec3], vel: wp.array[wp.vec3]): i wp.tid() position pos[i] dist_sq wp.length_sq(position) 0.01 # softened distance acc -1000.0 / dist_sq * wp.normalize(position) # gravitational pull toward origin vel[i] vel[i] acc * dt pos[i] pos[i] vel[i] * dt rng np.random.default_rng(42) positions wp.array(rng.normal(size(num_particles, 3)), dtypewp.vec3) velocities wp.array(rng.normal(size(num_particles, 3)), dtypewp.vec3) for _ in range(100): wp.launch(gravity_step, dimnum_particles, inputs[positions, velocities]) print(positions.numpy())这段代码浓缩了 Warp 的四个核心要素逐一拆解如下。1.wp.kernel声明一个内核wp.kernel装饰器把一个 Python 函数注册为 Warp 内核。由源码中的 kernel() 定义 可知它的约束很明确函数的所有参数必须带类型注解且函数不能有返回值。这里的参数类型wp.array[wp.vec3]声明了元素为三维向量vec3的数组wp.vec3是 Warp 内置的数学向量类型定义在 warp/_src/math.py 与 warp/native/vec.h 对应实现中。除基本用法外kernel()还支持若干控制编译行为的可选参数可参见 源码 docstring参数作用name自定义内核注册与代码生成使用的 key必须是合法的 C 标识符enable_backward设为False时跳过反向梯度pass 的生成可减少编译量launch_boundsCUDA__launch_bounds__属性可为 intmaxThreadsPerBlock或(maxThreadsPerBlock, minBlocksPerMultiprocessor)元组仅对 CUDA 生效且wp.launch的block_dim不能超过其maxThreadsPerBlockcuda_max_registersCUDA__maxnreg__寄存器上限不能与launch_bounds混用enable_cuda_smem_spillingCUDA Toolkit 13 时允许用共享内存做寄存器溢出cluster_dimCUDA Thread Block Cluster 大小1D CTA 数量≤16仅对计算能力 ≥9.0 的 Hopper 及更新架构生效module/moduleunique内核所属的编译模块传unique时为其单独创建新模块module_options模块级编译选项如fast_math、mode、max_unroll、deterministic需搭配moduleunique使用entry_point_abi实验性入口 ABIexternal_constant_params仅支持 CUDAgrid_stride是否使用 grid-stride 循环配合launch的max_blocks2.wp.tid()获取线程索引wp.tid()返回当前线程的全局索引是内核内并行度分配的基础。与 CUDA 的blockIdx/threadIdx组合不同Warp 将启动维度抽象为一维或最多四维的线程网格tid()直接给出该网格中的逻辑索引写法更简洁。启动维度大于一维时wp.tid()返回多分量索引例如二维数组遍历时返回(i, j)可参考 kernel() docstring 中的示例。3.wp.arrayGPU 就绪的数据容器wp.array是 Warp 的核心数据容器。示例中用wp.array(..., dtypewp.vec3)从 NumPy 数组创建既可以通过positions.numpy()拷回 CPU也能与 NumPy/PyTorch/JAX 生态互操作仓库提供了 warp/_src/torch.py、warp/_src/jax/、warp/_src/paddle.py 等互操作层。数组元素类型覆盖vec3、float、int及自定义结构体等。4.wp.launch异步启动内核wp.launch在内核定义与数据就绪后执行启动。由 launch() 源码 可见其完整签名与关键行为wp.launch( kernel, # 被 wp.kernel 装饰的内核 dim, # 线程数int 或最多 4 维的整数序列 inputs[...], # 输入参数可选 outputs[...], # 输出参数可选 deviceNone, # 目标设备cpu 或 CUDA 设备 streamNone, # 指定流指定后 device 由流的设备决定 block_dim256, # 每块线程数CPU 上强制为 1 max_blocks0, # CUDA 线程块上限需要 grid-stride 循环配合 ... )值得注意的实现细节启动是异步的与调用它的 Python 线程解耦源码 docstring 明确说明Kernel launches are asynchronous with respect to the calling Python thread。block_dim默认 256当目标设备为 CPU 时源码会强制将block_dim置为 1即 CPU 后端按单线程块串行展开。启动前会校验kernel必须是wp.kernel装饰过的Kernel对象并检查实参数量与内核签名是否匹配否则抛出RuntimeErrorwarp/_src/context.py#L11335-L11336。还可以通过adj_inputs/adj_outputs/adjoint/record_tape参与自动微分与 warp/_src/tape.py 的梯度记录这是 Warp 内核可微分的入口。此外示例中用dist_sq wp.length_sq(position) 0.01做了引力软化softening避免粒子经过原点时出现无穷大加速度是粒子模拟中常见的数值稳定技巧实际复用时可根据物理尺度调整软化半径。安装平台支持与多种安装途径环境要求与发布平台Warp 需要Python 3.10 或更高版本这一点在 pyproject.toml 中以requires-python 3.10显式声明README 与包元数据保持一致。Warp 在 PyPI 上以warp-lang包名发布覆盖以下平台平台架构CPU 执行CUDA 加速Windowsx86-64✓✓需 NVIDIA GPU 与驱动Linuxx86-64 / AArch64✓✓需 NVIDIA GPU 与驱动macOSApple Silicon✓✗不支持 Metal 加速注意CUDA 加速需要受支持的 NVIDIA GPU 与驱动macOS 版本仅支持 CPU 执行。基础安装pip install warp-lang这是最简单的方式。若想运行仓库自带示例以及使用 USD 相关功能安装带examplesextra 的版本pip install warp-lang[examples]examplesextra 具体会带入哪些依赖可以从 pyproject.toml 的 optional-dependencies 段 查到包括matplotlib、pillow、psutil、pyglet、blosc以及usd-core用于 USD 文件读写等。仓库还定义了其他可选的安装组合warp-lang[torch-cu12]warp-lang[examples]torch2.7.0面向 CUDA 12 的 PyTorch 集成warp-lang[torch-cu13]warp-lang[examples]torch2.9.0面向 CUDA 13warp-lang[dev]含 examples、docs 及coverage、ml-dtypes、nvtx等开发依赖适用于参与贡献与本地开发。一个值得注意的细节在 Linux aarch64 系统如 NVIDIA DGX Spark上由于usd-core没有对应平台的 wheel[examples]extra 会自动安装usd-exchange作为替代实现保持功能对齐。从源码构建如需夜间构建、conda 环境、CUDA 13 构建或从源码构建可参考仓库中的构建脚本与文档build_lib.py 是核心运行库的“离线”构建脚本面向 CI 与开发者工作流需要 CUDA Toolkit 等工具链不是用户运行时的一部分build_llvm.py 负责构建 Warp 依赖的 LLVM 组件详细的安装指南见 docs/user_guide/installation.rst。从源码构建时构建过程会自动下载 NVIDIA libmathdx而 PyPI 预编译包已将 libmathdx 静态链接进库文件。无论哪种方式libmathdx 都受 NVIDIA Software License Agreement 约束许可文本见 licenses/libmathdx-LICENSE.txt。仓库还提供了基于容器的一键构建环境 docker/warp-builder/ 与 CUDA 测试环境 docker/warp-cpp-test-env/。运行官方示例从命令行到可视化安装warp-lang[examples]后即可运行仓库 warp/examples/ 下的各类示例覆盖物理模拟、几何处理、优化以及基于 tile 的 GPU 编程。命令行运行python -m warp.examples.example_subdir.example例如运行核心示例中的离散元模拟dem与流固耦合sphpython -m warp.examples.core.example_dem python -m warp.examples.core.example_sph绝大多数示例既可在 CPU 上运行也可在支持 CUDA 的设备上运行少数示例必须在 CUDA 设备上运行这类脚本会在文件顶部明确标注README 明确说明了这一点运行前值得先看一眼脚本头部注释。部分示例会在当前工作目录生成包含时间采样动画的 USD 文件可用 Pixar UsdView、Blender 或其他兼容 USD 的查看器打开。浏览示例源码想快速浏览示例源码目录可以运行仓库自带的浏览工具python -m warp.examples.browse它位于 warp/examples/browse.py会打开示例所在目录方便按需阅读源码。示例全景仓库示例按领域组织为四个主要子目录对应 warp/examples/ 下的实际结构core核心计算——warp/examples/core/离散元 example_dem.py、流固耦合 example_sph.py、图捕获 example_graph_capture.py、marching cubes example_marching_cubes.py、网格 example_mesh.py、NVDB 稀疏体积 example_nvdb.py、光线投射 example_raycast.py、光线步进 example_raymarch.py、网格采样 example_sample_mesh.py、波动 example_wave.py、PyTorch 互操作 example_torch.py、周期盒中的二维不可压缩湍流 example_fft_poisson_navier_stokes_2d.py、三维 FDTD Luneburg 透镜 example_fdtd_3d.py 等。fem有限元法——warp/examples/fem/三维扩散 example_diffusion_3d.py、混合弹性 example_mixed_elasticity.py、APIC 流体 example_apic_fluid.py、流线 example_streamlines.py、畸变能 example_distortion_energy.py、Taylor-Green 涡 example_taylor_green.py、Kelvin-Helmholtz 不稳定性 example_kelvin_helmholtz.py、静磁学 example_magnetostatics.py、自适应网格 example_adaptive_grid.py、非协调接触 example_nonconforming_contact.py、Darcy level-set 优化 example_darcy_ls_optimization.py、弹性形状优化 example_elastic_shape_optimization.py 等。optim优化——warp/examples/optim/可微光线追踪 example_diffray.py、带检查点的流体优化 example_fluid_checkpoint.py、粒子排斥优化 example_particle_repulsion.py、Navier-Stokes 扰动优化 example_navier_stokes_perturbation.py。tile基于 tile 的 GPU 编程——warp/examples/tile/多层感知机 example_tile_mlp.py、N 体问题 example_tile_nbody.py、多通道高斯过程 example_tile_mcgp.py。此外还有 warp/examples/distributed/分布式、warp/examples/interop/框架互操作、warp/examples/cpp/C/CUDA 工作流与 warp/examples/benchmarks/性能基准含与 PyTorch、JAX、Taichi、Numba、CuPy 等的对比脚本等目录可按需探索。深入 Warp 的推荐路径从 README 的 Quick Start 出发可以沿以下仓库内资源逐步深入编程模型docs/user_guide/programming_model/code_generation.rst 讲解代码生成机制docs/user_guide/programming_model/generics.rst 讲解泛型内核docs/user_guide/programming_model/tiles.rst 讲解 tile 编程docs/user_guide/basics.rst 覆盖语言基础。可微分性与自动微分docs/user_guide/differentiability.rst 与 docs/user_guide/interoperability/pytorch.rst、docs/user_guide/interoperability/jax.rst 分别说明梯度与框架集成。运行时与性能docs/user_guide/runtime.rst、docs/user_guide/execution_and_performance/ 下的并发、内存管理、剖析与启动优化文档。API 参考docs/api_reference/ 中的warp.rst、warp_types.rst、warp_autograd.rst等覆盖完整 API语言内置函数参考见 docs/language_reference/builtins.rst。设计文档design/ 目录记录了确定性执行、统一日志、可插拔分配器、外部编译扩展等专题设计。变更记录CHANGELOG.md 与 changelog/ 目录收录了各版本新增、修复与行为变更。如果参与开发或研究引用请参阅 CONTRIBUTING.md贡献指南、PUBLICATIONS.md使用 Warp 的学术出版物列表与 CITATION.cff引用信息。许可证与第三方组件Warp 采用Apache License 2.0发布完整许可文本见 LICENSE.md。需要注意的是使用过程中会自动下载并安装额外的第三方开源软件使用前应审查这些开源项目的许可条款licenses/ 目录集中存放了各第三方组件的许可文本例如 CUDA、LLVM、libmathdx、nanovdb、USDL 等。README 中的 NOTICE 声明强调这些自动获取的外部材料不随软件分发受各自条款约束使用者需自行负责审查并验证其安全性、完整性及适用性。【免费下载链接】warpA Python framework for GPU-accelerated simulation, robotics, and machine learning.项目地址: https://gitcode.com/GitHub_Trending/warp/warp创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表