tessl/pypi-warp-lang

A Python framework for high-performance simulation and graphics programming that JIT compiles Python functions to efficient GPU/CPU kernel code.

Overview

Eval results

Files

Types and Arrays

Name: tessl/pypi-warp-lang
Author: tessl

Warp provides a comprehensive type system including primitive types, vectors, matrices, quaternions, transforms, and multi-dimensional arrays with device-aware memory management. This type system enables efficient GPU/CPU computation while maintaining Python accessibility.

Capabilities

Primitive Types

Basic numeric types optimized for GPU computation with explicit bit-width specification.

# Integer types
int8 = typing.Type[int]     # 8-bit signed integer
uint8 = typing.Type[int]    # 8-bit unsigned integer  
int16 = typing.Type[int]    # 16-bit signed integer
uint16 = typing.Type[int]   # 16-bit unsigned integer
int32 = typing.Type[int]    # 32-bit signed integer
uint32 = typing.Type[int]   # 32-bit unsigned integer
int64 = typing.Type[int]    # 64-bit signed integer
uint64 = typing.Type[int]   # 64-bit unsigned integer

# Floating point types
float16 = typing.Type[float]  # 16-bit half precision
float32 = typing.Type[float]  # 32-bit single precision
float64 = typing.Type[float]  # 64-bit double precision

# Boolean type
bool = typing.Type[bool]

Vector Types

Fixed-size vector types with component-wise operations for spatial computing.

# 2D vectors with different component types
vec2 = typing.Type[Vector2]      # float32 components
vec2b = typing.Type[Vector2]     # int8 components
vec2ub = typing.Type[Vector2]    # uint8 components  
vec2s = typing.Type[Vector2]     # int16 components
vec2us = typing.Type[Vector2]    # uint16 components
vec2i = typing.Type[Vector2]     # int32 components
vec2ui = typing.Type[Vector2]    # uint32 components
vec2l = typing.Type[Vector2]     # int64 components
vec2ul = typing.Type[Vector2]    # uint64 components
vec2h = typing.Type[Vector2]     # float16 components
vec2f = typing.Type[Vector2]     # float32 components
vec2d = typing.Type[Vector2]     # float64 components

# 3D vectors with different component types
vec3 = typing.Type[Vector3]      # float32 components
vec3b = typing.Type[Vector3]     # int8 components
vec3ub = typing.Type[Vector3]    # uint8 components
vec3s = typing.Type[Vector3]     # int16 components
vec3us = typing.Type[Vector3]    # uint16 components
vec3i = typing.Type[Vector3]     # int32 components
vec3ui = typing.Type[Vector3]    # uint32 components
vec3l = typing.Type[Vector3]     # int64 components
vec3ul = typing.Type[Vector3]    # uint64 components
vec3h = typing.Type[Vector3]     # float16 components
vec3f = typing.Type[Vector3]     # float32 components
vec3d = typing.Type[Vector3]     # float64 components

# 4D vectors with different component types  
vec4 = typing.Type[Vector4]      # float32 components
vec4b = typing.Type[Vector4]     # int8 components
vec4ub = typing.Type[Vector4]    # uint8 components
vec4s = typing.Type[Vector4]     # int16 components
vec4us = typing.Type[Vector4]    # uint16 components
vec4i = typing.Type[Vector4]     # int32 components
vec4ui = typing.Type[Vector4]    # uint32 components
vec4l = typing.Type[Vector4]     # int64 components
vec4ul = typing.Type[Vector4]    # uint64 components
vec4h = typing.Type[Vector4]     # float16 components
vec4f = typing.Type[Vector4]     # float32 components
vec4d = typing.Type[Vector4]     # float64 components

Matrix Types

Square matrix types for linear algebra and transformations.

# 2x2 matrices
mat22 = typing.Type[Matrix22]    # float32 components
mat22h = typing.Type[Matrix22]   # float16 components  
mat22f = typing.Type[Matrix22]   # float32 components
mat22d = typing.Type[Matrix22]   # float64 components

# 3x3 matrices
mat33 = typing.Type[Matrix33]    # float32 components
mat33h = typing.Type[Matrix33]   # float16 components
mat33f = typing.Type[Matrix33]   # float32 components  
mat33d = typing.Type[Matrix33]   # float64 components

# 4x4 matrices
mat44 = typing.Type[Matrix44]    # float32 components
mat44h = typing.Type[Matrix44]   # float16 components
mat44f = typing.Type[Matrix44]   # float32 components
mat44d = typing.Type[Matrix44]   # float64 components

Quaternion and Transform Types

Specialized types for 3D rotations and transformations.

# Quaternions for rotations
quat = typing.Type[Quaternion]    # float32 components
quath = typing.Type[Quaternion]   # float16 components
quatf = typing.Type[Quaternion]   # float32 components  
quatd = typing.Type[Quaternion]   # float64 components

# 3D transforms (rotation + translation)
transform = typing.Type[Transform]    # float32 components
transformh = typing.Type[Transform]   # float16 components
transformf = typing.Type[Transform]   # float32 components
transformd = typing.Type[Transform]   # float64 components

Spatial Mathematics Types

Specialized types for spatial mechanics and robotics.

# 6D spatial vectors (angular + linear velocity/force)
spatial_vector = typing.Type[SpatialVector]    # float32 components
spatial_vectorh = typing.Type[SpatialVector]   # float16 components
spatial_vectorf = typing.Type[SpatialVector]   # float32 components
spatial_vectord = typing.Type[SpatialVector]   # float64 components

# 6x6 spatial matrices (inertia, compliance)
spatial_matrix = typing.Type[SpatialMatrix]    # float32 components
spatial_matrixh = typing.Type[SpatialMatrix]   # float16 components
spatial_matrixf = typing.Type[SpatialMatrix]   # float32 components
spatial_matrixd = typing.Type[SpatialMatrix]   # float64 components

Array Types

Multi-dimensional arrays with device-aware memory management.

class array:
    """
    Multi-dimensional array with device-aware memory management.
    
    Attributes:
        shape: Tuple of dimension sizes
        dtype: Element data type
        device: Device where array is stored
        size: Total number of elements
    """
    
    def __init__(self, data=None, dtype=float32, shape=None, device=None):
        """
        Create array from data or with specified shape.
        
        Args:
            data: Initial data (numpy array, list, or scalar)
            dtype: Element type (warp type or numpy dtype)
            shape: Array dimensions if data is None
            device: Target device for array storage
        """
    
    def numpy(self) -> np.ndarray:
        """Convert to NumPy array (copies from device if needed)."""
    
    def __getitem__(self, key):
        """Get array element or slice."""
    
    def __setitem__(self, key, value):
        """Set array element or slice."""

# Type aliases for common array dimensions
array1d = array  # 1D array
array2d = array  # 2D array
array3d = array  # 3D array
array4d = array  # 4D array

# Indexed arrays for sparse data access
indexedarray = typing.Type[IndexedArray]
indexedarray1d = typing.Type[IndexedArray]
indexedarray2d = typing.Type[IndexedArray]
indexedarray3d = typing.Type[IndexedArray]
indexedarray4d = typing.Type[IndexedArray]

# Fabric arrays for distributed computing
fabricarray = typing.Type[FabricArray]
fabricarrayarray = typing.Type[FabricArrayArray]
indexedfabricarray = typing.Type[IndexedFabricArray]
indexedfabricarrayarray = typing.Type[IndexedFabricArrayArray]

# Tiled arrays for blocked algorithms
tile = typing.Type[Tile]

Geometry Types

Spatial data structures for collision detection, ray tracing, and geometric queries.

class Mesh:
    """Triangle mesh for collision detection and rendering."""
    
    def __init__(self, vertices: array, indices: array):
        """
        Create mesh from vertex positions and triangle indices.
        
        Args:
            vertices: Array of 3D vertex positions
            indices: Array of triangle vertex indices
        """

class Volume:
    """3D volume for sampling and level sets."""
    
    def __init__(self, data: array, transform: transform = None):
        """
        Create volume from 3D scalar field data.
        
        Args:
            data: 3D array of scalar values
            transform: World-space transformation
        """

class Bvh:
    """Bounding Volume Hierarchy for fast spatial queries."""
    
    def __init__(self, mesh: Mesh):
        """Build BVH from triangle mesh."""

class HashGrid:
    """Spatial hash grid for neighbor finding."""
    
    def __init__(self, dim_x: int, dim_y: int, dim_z: int):
        """Create hash grid with specified dimensions."""

class MarchingCubes:
    """Marching cubes for isosurface extraction."""
    
    def __init__(self):
        """Create marching cubes instance."""

Query Types

Spatial query objects for geometric computations.

class BvhQuery:
    """Query object for BVH spatial searches."""

class HashGridQuery:
    """Query object for hash grid neighbor searches."""

class MeshQueryAABB:
    """Axis-aligned bounding box mesh query."""

class MeshQueryPoint:
    """Point-to-mesh distance query."""

class MeshQueryRay:
    """Ray-mesh intersection query."""

Utility Functions

Functions for creating and manipulating arrays and types.

def constant(value, dtype: type = None):
    """Create compile-time constant value."""

def from_ptr(ptr: int, dtype: type, shape: tuple) -> array:
    """Create array from raw memory pointer."""

def from_ipc_handle(handle, dtype: type, shape: tuple) -> array:
    """Create array from CUDA IPC memory handle."""

Type Annotations

Special types for kernel parameter annotations and generic programming.

Int = typing.TypeVar('Int')        # Any integer type
Float = typing.TypeVar('Float')    # Any floating point type  
Scalar = typing.TypeVar('Scalar')  # Any numeric type

NumPy Interoperability

Functions for converting between Warp and NumPy types.

def dtype_from_numpy(np_dtype) -> type:
    """Convert NumPy dtype to Warp type."""

def dtype_to_numpy(wp_dtype: type):
    """Convert Warp type to NumPy dtype."""

Usage Examples

Vector Operations

import warp as wp

# Create 3D vectors
pos = wp.vec3(1.0, 2.0, 3.0)
vel = wp.vec3(0.1, 0.2, 0.3)

# Vector operations available in kernels
@wp.kernel
def update_positions(positions: wp.array(dtype=wp.vec3),
                    velocities: wp.array(dtype=wp.vec3),
                    dt: float):
    i = wp.tid()
    positions[i] = positions[i] + velocities[i] * dt

Matrix Transformations

# 4x4 transformation matrix
transform_matrix = wp.mat44(
    1.0, 0.0, 0.0, 0.0,
    0.0, 1.0, 0.0, 0.0, 
    0.0, 0.0, 1.0, 0.0,
    0.0, 0.0, 0.0, 1.0
)

# Quaternion rotation
rotation = wp.quat(0.0, 0.0, 0.0, 1.0)  # Identity rotation

Array Creation and Access

import numpy as np

# Create arrays on different devices
cpu_array = wp.zeros(1000, dtype=wp.float32, device="cpu")
gpu_array = wp.ones((100, 100), dtype=wp.vec3, device="cuda:0")

# Convert from NumPy
np_data = np.random.rand(1000).astype(np.float32)
wp_array = wp.from_numpy(np_data, device="cuda:0")

# Access data
result = wp_array.numpy()  # Copy back to NumPy

Types

# Vector component access (available in kernels)
class Vector2:
    x: float
    y: float

class Vector3:
    x: float
    y: float
    z: float

class Vector4:
    x: float
    y: float
    z: float
    w: float

# Matrix types
class Matrix22:
    def __getitem__(self, key: tuple[int, int]) -> float: ...

class Matrix33:
    def __getitem__(self, key: tuple[int, int]) -> float: ...

class Matrix44:
    def __getitem__(self, key: tuple[int, int]) -> float: ...

# Quaternion type
class Quaternion:
    x: float  # i component
    y: float  # j component  
    z: float  # k component
    w: float  # real component

# Transform type (rotation + translation)
class Transform:
    p: vec3  # translation
    q: quat  # rotation

Install with Tessl CLI