NVIDIA / warp

A Python framework for high performance GPU simulation and graphics

Home Page:https://nvidia.github.io/warp/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

wp.struct: from numpy and back

jc211 opened this issue · comments

Hi,
I am wondering how to move wp.arrays of type my_struct back and forth between numpy and cuda. At the moment, only one direction seems to work:

@wp.struct
class foo:
    temp: int

a_d = wp.zeros((10,), dtype=my_struct) # works
a = a_d.numpy() # works
a_d = wp.from_numpy(a, dtype=my_struct) # does not work

The error is:

warp/warp/types.py", line 1325, in _init_from_data
    raise RuntimeError(f"Error while trying to construct Warp array from a sequence of Warp structs: {e}")
RuntimeError: Error while trying to construct Warp array from a sequence of Warp structs: 'numpy.void' object has no attribute '__ctype__

Am I doing something wrong?

It could be done like below:

@wp.struct
class foo:
    temp: int

a_d = wp.zeros((10,), dtype=my_struct) 
a = a_d.to("cpu") 
a_numpy = np.asarray(a, dtype=my_struct.numpy_dtype())
a_numpy["temp"] = 3

Although it would be cleaner if from_numpy() worked on structs

Hi @jc211, I agree that wp.from_numpy() should handle this gracefully. We'll get that fixed shortly. In the meantime, please consider using this helper function:

def from_numpy_structs(a, dtype, device=None):
    # wrap numpy array as a host array using the pointer
    a_h = wp.array(ptr=a.ctypes.data, dtype=dtype, shape=a.shape, device="cpu", copy=False, owner=False)
    # resolve the device
    device = wp.get_device(device)
    if device.is_cpu:
        # keep a back reference to prevent the numpy array from getting deallocated
        a_h._ref = a
        return a_h
    else:
        # copy to device
        return a_h.to(device)

Your code would change like this:

a_d = wp.zeros((10,), dtype=my_struct) # works
a = a_d.numpy() # works
a_d = from_numpy_structs(a, my_struct)  # workaround

Hi @jc211 , this has been addressed in 46f1bc7.

Cheers,
Miles