Zero-Copy

The Python binding pyAMReX bridges the compute in AMReX block-structured codes and data science. As such, it includes zero-copy GPU data access for AI/ML, in situ analysis, application coupling by implementing standardized data interfaces.

CPU: NumPy

zero-copy read and write access. CPU as well as managed memory CPU/GPU.

Call .to_numpy() on data objects of pyAMReX. See the optional arguments of this API.

Writing to the created NumPy array will also modify the underlying AMReX memory.

GPU: CuPy

GPU zero-copy read and write access on NVIDIA (CUDA) and AMD (ROCm) GPUs.

Call .to_cupy() on data objects of pyAMReX. See the optional arguments of this API.

Writing to the created CuPy array will also modify the underlying AMReX memory.

GPU: dpnp

GPU zero-copy read and write access on Intel (SYCL) GPUs, exchanged via DLPack.

Call .to_dpnp() on data objects of pyAMReX. See the optional arguments of this API.

Writing to the created dpnp array will also modify the underlying AMReX memory.

CPU/GPU Agnostic Code: NumPy/CuPy/dpnp

The previous examples can be written in CPU/GPU agnostics manner. Either using NumPy (np), CuPy (cp) or dpnp (dp), we provide a common short-hand abbreviation named xp .

Call .to_xp() on data objects of pyAMReX. It returns NumPy arrays on CPU builds, CuPy arrays on CUDA/ROCm builds and dpnp arrays on SYCL builds. See the optional arguments of this API.

Writing to the created NumPy/CuPy/dpnp array will also modify the underlying AMReX memory.

The matching array library itself is available as xp on the pyAMReX module, for code that has to call into it rather than just hold its arrays:

import amrex.space3d as amr

arr = mfab.array(mfi).to_xp()
amr.xp.sin(arr, out=arr)  # numpy, cupy or dpnp, matching the build

It is resolved on first access, so importing pyAMReX does not import a GPU array library.

GPU: numba

GPU zero-copy read and write access.

After from numba import cuda, create a zero-copy tensor on a GPU array via marr_numba = cuda.as_cuda_array(marr).

Writing to the created numba array will also modify the underlying AMReX memory.

AI/ML: pyTorch

CPU and GPU zero-copy read and write access.

Create a zero-copy tensor on a GPU array via torch.as_tensor(amrex_array_here, device="cuda") or torch.from_dlpack(amrex_array_here).

Writing to the created PyTorch tensor will also modify the underlying AMReX memory.

Everything Else: DLPack

Frameworks not covered above can exchange data with pyAMReX through the standardized DLPack protocol, on CPU as well as on CUDA, ROCm and SYCL GPUs.

Data objects of pyAMReX implement __dlpack__ and __dlpack_device__ and can be consumed by, e.g., numpy.from_dlpack, cupy.from_dlpack, dpnp.from_dlpack, torch.from_dlpack or jax.dlpack. Both DLPack 1.x (“versioned”) and legacy consumers are supported, including device-to-host transfers via from_dlpack(..., device="cpu").

A consuming NumPy, CuPy, dpnp, PyTorch or other DLPack array must be released before calling amrex.finalize(). Finalization runs Python’s garbage collector and raises RuntimeError if a DLPack managed tensor is still alive, because it can retain an AMReX object or arena allocation after runtime teardown. Delete all consuming arrays and capsules, then retry finalization.

An exception is the deprecated particle ArrayOfStructs: DLPack cannot describe its record (struct) element type. Use its structured __array_interface__/__cuda_array_interface__ or the per-component particle struct-of-arrays data instead.