DOCSEARCHbyAlgolia
Submit Docs

NVIDIA Warp Documentation

repository·main·Indexed 27 days ago

https://github.com/nvidia/warp

A Python framework (warp-lang) for high-performance, GPU-accelerated simulation, robotics, and machine learning. It uses JIT compilation to convert Python functions into efficient CPU/GPU kernels and supports differentiability for integration with ML frameworks like PyTorch, JAX, and Paddle.

Tokens
109.7K
Snippets
245
Records
532
Agent score
92%

What's inside NVIDIA Warp

  1. Overview of NVIDIA Warp

    main

    NVIDIA Warp is a Python framework designed for GPU-accelerated simulation, robotics, and machine learning. It allows developers to write regular Python functions and JIT (Just-In-Time) compile them into efficient kernels that execute on either the CPU or GPU.

    Key features include:

    • Differentiable kernels suitable for machine learning pipelines (compatible with PyTorch, JAX, and Paddle).
    • Primitives for physics simulation, robotics, and geometry processing.

  2. Overview of LLVM SDK Build Methods

    main

    Warp uses a custom, size-optimized LLVM/Clang flavor (clang-only static libraries with the NVPTX backend) as its CPU-JIT compiler. There are three existing ways to build this SDK:

    1. build_llvm.py: A plain CMake driver used as a user-facing fallback via build_lib.py --build-llvm.
    2. docker/warp-builder/Dockerfile: Uses manylinux containers to bake /opt/llvm into builder images.
    3. Conan Recipe (clang-warp): An NVIDIA-internal recipe that supports per-platform profiles, cross-building for windows-aarch64, size optimization, and library pruning.

    The project is moving towards a public, standalone Conan recipe located in tools/llvm/ that can be executed on GitHub Actions runners.

  3. Understand the NVIDIA Warp programming model

    main

    Warp is a Python framework for writing kernels that run on CPUs and NVIDIA GPUs. It allows authors to express parallel work using logical thread indices, array operations, and kernel launches.

    Key features include:

    • Compilation: Lowers typed Python kernel code to generated C++ (for CPU via LLVM/Clang) or CUDA C++ (for GPU via NVRTC).
    • Automatic Differentiation: Generates reverse-mode automatic differentiation code for supported kernels and functions.
    • Spatial API: Includes GPU-accelerated BVHs, hash grids, triangle meshes, and sparse volumes with query primitives.
    • Simulation Toolkit: Includes sparse linear algebra and a finite element method (FEM) toolkit for building solvers.
    • Interoperability: Designed to work with NumPy, PyTorch, and JAX using array interfaces, DLPack, and dedicated converters to enable zero-copy data sharing where supported.

  4. Use warp.fem.polynomial for Finite Element Method (FEM) polynomial operations

    main
    The warp.fem.polynomial module provides utilities for polynomial-based operations within the Finite Element Method (FEM) framework in NVIDIA Warp. It includes functions for calculating Lagrange scales and performing 1D quadrature, which are essential for integrating polynomial basis functions over elements.

  6. Understand Warp versioning and release types

    main

    Warp uses an X.Y.Z versioning format that does not strictly follow Semantic Versioning. Understanding the components helps you anticipate changes:

    • X (Marketing Number): Reserved for major reworks causing disruptive incompatibility.
    • Y (Feature Release): Published monthly. These are the only releases that may contain new features, deprecations, breaking changes, or removals.
    • Z (Bugfix Release): Issued ad-hoc for important issues. These never contain new features, deprecations, or removals.

    Prerelease formats:

    • X.Y.Z.dev0: Development builds from the main branch.
    • X.Y.ZrcN: Release candidates (e.g., 1.10.0rc1) used for QA.
    • X.Y.Z.devYYYYMMDD: Nightly builds published on the NVIDIA PyPI index.

  7. Use warp.fem.linalg for finite element linear algebra operations

    main
    The warp.fem.linalg module provides a collection of specialized linear algebra functions designed for finite element method (FEM) applications. These functions include operations for matrix decompositions (QR, Hessenberg), eigenvalue problems (symmetric, tridiagonal), and matrix manipulations (symmetric/skew parts, AXPY).

  8. Use shape functions in warp.fem.space.shape

    main
    The warp.fem.space.shape module provides various shape function implementations for Finite Element Method (FEM) applications. These functions are categorized by the geometric primitive they operate on (e.g., Cube, Square, Tetrahedron, Triangle) and the specific mathematical type of the shape function (e.g., BSpline, NedelecFirstKind, RaviartThomas, Serendipity, Polynomial).

  9. Hardware-Coherent Cross-Device Memory Access Overview

    main
    Warp is implementing support for hardware-coherent cross-device memory access. Historically, Warp enforced a strict rule that all array arguments passed to wp.launch() must reside on the same device as the kernel launch target. On systems with unified memory architectures (like Grace C2C, Jetson Thor, or HMM-enabled Linux systems), this restriction is being relaxed to allow the GPU to directly access CPU memory without requiring explicit wp.copy() or .to(device) calls. This allows developers to use ordinary system allocations (e.g., malloc, mmap) directly in GPU kernels on supported hardware.

  10. Use the warp.sparse module for Block Sparse Row (BSR) operations

    main
    The warp.sparse module provides tools for working with Block Sparse Row (BSR) matrices. It includes support for matrix creation from triplets, identity matrices, zero matrices, and various linear algebra operations such as matrix-vector multiplication (bsr_mv), matrix-matrix multiplication (bsr_mm), and scaling. It also provides utilities for managing block indices and compression.

  11. Understand CPU vs CUDA execution in Warp

    main

    Warp supports both CPU and CUDA backends, but they behave differently:

    • CPU Execution: Kernel launches are serial and synchronous. Note that for Tile kernels, the effective block_dim is 1 on the CPU backend.
    • CUDA Execution: Kernel launches run many threads in parallel and are generally asynchronous with respect to Python.

    Both backends use the same kernel language, but performance and concurrency characteristics differ.

  12. Integrate C++/CUDA with Warp workflows

    main

    While Warp is primarily authored in Python, it provides several workflows for non-Python integration. You can use Warp to:

    • Insert native C++/CUDA snippets directly into generated Warp kernels.
    • Perform Ahead-of-Time (AOT) compilation of Warp kernels into source code, PTX, or CUBIN files.
    • Load generated Warp binaries or source code from a native CUDA C++ application.
    • Serialize captured Warp work and replay it from a C++ application without requiring a Python runtime.