Open source at last: Warp, Nvidia's Python framework for CUDA

Nvidia has placed Warp under an Apache 2 license. The Python framework is used for performance-hungry physical simulations, data generation and spatial computing. It compiles Python functions just in time into kernel code that runs on x86 CPUs and CUDA GPUs.

Nvidia itself remains silent about the background to the license change. However, there had previously been criticism in the community of the original proprietary license, which disclosed the sources but reserved rights of use. The new license is now recognized as open source.

Comparable with native CUDA code

Warp provides special primitives for physical simulations, sensors, robotics or geometric processing. This also includes higher-level geometric data structures such as meshes, hash grids and sparse volumes for use in the GPU. In contrast to tensor-based programming, Warp offers implicit, developer-controlled kernel and threat management, native conditional logic as well as sparse scatter and gather for vector processing.

The Warp kernels play back information in reverse mode for use in frameworks such as PyTorch, JAX, PhysicsNeMo, Nvidia Omniverse or Paddle. Use in pipelines for machine learning is also planned. Nvidia praises Warp on the project's website, stating that the “simulation speed is equivalent to native CUDA code with the convenience and developer productivity of Python”.

Warp runs on the x86-64 or ARMv8 platforms under Windows, Linux, and macOS. GPU support requires CUDA GPUs, drivers, and graphics cards of at least the GeForce GTX 9xx series. For Python, Nvidia recommends version 3.9 and higher. The binaries on PyPI are built with CUDA 12 and require CUDA drivers from 525.60.13 (Linux x86-64) or 528.33 (Windows x86-64). The installation is done with:

pip install warp-lang

In the Warp repo there is a code example that calculates the length of random 3D vectors:

import warp as wp
import numpy as np

num_points = 1024

@wp.kernel
def length(points: wp.array(dtype=wp.vec3),
           lengths: wp.array(dtype=float)):

    # thread index
    tid = wp.tid()
    
    # compute distance of each point from origin
    lengths[tid] = wp.length(points[tid])


# allocate an array of 3d points
points = wp.array(np.random.rand(num_points, 3), dtype=wp.vec3)
lengths = wp.zeros(num_points, dtype=float)

# launch kernel
wp.launch(kernel=length,
          dim=len(points),
          inputs=[points, lengths])

print(lengths)

Further information can be found in a few Jupyter notebooks, in the documentation and in code examples.

Lesen Sie auch

data2day 2025: Vorschläge zur Data-Science-Konferenz noch bis 19. Mai einreichen

(who)