Expose WarpX data to Python, on GPU

[RemiLehe]

Right now, WarpX exposes its data to Python, by converting it to numpy arrays:

WarpX/Python/pywarpx/_libwarpx.py

Line 135 in 1f8ab0e

return np.frombuffer(buf, dtype=dtype, count=size)

However, this would not be a good solution when running on GPU. Instead we should expose the data as a GPU-aware object (e.g. cupy array)

[ax3l]

cupy

Comparison table has no frombuffer equivalent: https://docs.cupy.dev/en/stable/reference/comparison.html

Chainer mentions a cupy.creation.from_data: https://docs.chainer.org/en/v1.3.1/_modules/cupy/creation/from_data.html
Code

cupy.ndarray constructors not obvious to pass device pointers atm either: https://docs.cupy.dev/en/stable/reference/generated/cupy.ndarray.html#cupy.ndarray

We could explore if creating a numpy array with the pointer being a device pointer and calling cupy.asarray(numpy_array_w_device_ptr) from it works [ref] [more details] or if this triggers a copy and/or if this triggers ownership issues. There also seems to be a , copy=False option in array which is equivalent to asarray [ref].

Maybe we need to write our own constructor (could also consider contributing this upstream). I opened an issue to ask about details: cupy/cupy#4644

• definition
• a fallback wrapper

[ax3l]

Discussion moved to: AMReX-Codes/pyamrex#9

Looks like the community standardizes on the CUDA Array Interface (v3).

Ref.: cupy/cupy#4644 (comment)