Skip to content

Support wp.handle in framework and DLPack converters #2292

Description

@c0d1f1ed

Description

Support explicit dtype=wp.handle imports in wp.from_torch(), wp.from_jax(), wp.from_paddle(), and wp.from_dlpack() for compatible 64-bit integer storage. Support exporting handle arrays as 64-bit integer storage through DLPack and the converters that depend on it, including JAX and Paddle.

External frameworks do not need a custom handle dtype: the explicit Warp import dtype should mark the array as object references while preserving its underlying bits. Existing integer imports and exports should retain their behavior.

Context

#2134 recommends wp.handle for Mesh, Volume, BVH, and HashGrid IDs. wp.array() and wp.from_numpy() already accept dtype=wp.handle, but the other converters require a separate integer import and .view(wp.handle). This is easy to miss when sharing arrays of object IDs with a kernel annotated wp.array[wp.handle].

For example, the direct PyTorch import currently fails:

import torch
import warp as wp

# Illustrative ID storage; these zeros do not identify live resources.
source = torch.zeros(1, dtype=torch.int64)
handles = wp.from_torch(source, dtype=wp.handle)

RuntimeError: Cannot convert Torch type torch.int64 to Warp type <class 'warp._src.types.handle'>

The equivalent Paddle import also rejects wp.handle. Source inspection shows that DLPack's dtype compatibility and export mappings exclude wp.handle; JAX imports delegate to DLPack, and JAX/Paddle exports use it. PyTorch export uses array interfaces and should be checked separately rather than assuming it needs the same change.

Alternatives or Workarounds

Import a supported integer dtype and create an explicit handle view. Export an integer view supported by the target framework:

handles = wp.from_torch(source, dtype=wp.int64).view(wp.handle)
exported = wp.to_torch(handles.view(wp.int64))

Analogous signed 64-bit views work around the JAX, Paddle, and DLPack boundaries; JAX also requires 64-bit values to be enabled. These views reinterpret storage rather than numerically casting IDs.

Suggested validation

Cover supported CPU/CUDA routes, signed and unsigned 64-bit storage where the framework permits it, high-bit values, shape/stride compatibility, and storage ownership. Check imports, exports, and use in handle-typed kernel parameters without requiring intermediate views. Preserve the current integer workflows and clearly document any remaining framework restrictions.

This request concerns converter dtype support. Handle typing must continue to follow the existing resource lifetime and process-local ID rules; resource serialization remains separate work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

feature requestRequest for something to be added

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions