Implementation Details

For now, please see these presentations:

• A. Huebl et al., NREL virtual seminar, December 2022

• A. Huebl et al., “Exascale and ML Models for Accelerator Simulations”, 6th European Advanced Accelerator Concepts workshop (EAAC23), Isola d’Elba, Italy, Sep 17 – 23, 2023 DOI:10.5281/zenodo.8362549

Zero-Copy Data APIs

pyAMReX implements the following standardized data APIs:

• __array_interface__ (CPU)

• __cuda_array_interface__ (CUDA GPU)

• DLPack v1.1 via __dlpack__ and __dlpack_device__ (CPU and CUDA, ROCm & SYCL GPUs)

These APIs are automatically used when creating “views” (non-copy) numpy arrays, cupy arrays, dpnp arrays, PyTorch tensors, etc. from AMReX objects such as Array4 and particle arrays.

DLPack is implemented once in src/dlpack/DLPackHelpers.H (kwarg protocol, capsule creation, producer keep-alive, copy & stream synchronization) and bound per container class (Array4, BaseFab, PODVector, Vector, SmallMatrix) with a small adapter describing the exported tensor. An exception is ArrayOfStructs: DLPack cannot describe record (struct) element types, so it only implements the structured array interfaces above.