[WIP] Implement DLPack

[ax3l]

Add first-class support for zero-copy data exchange with ROCm and SYCL GPUs via DLPack interfaces.

Specs:

• https://dmlc.github.io/dlpack/latest/python_spec.html#implementation
• https://github.com/dmlc/dlpack/blob/v1.1/include/dlpack/dlpack.h

Note: we might want to implement a slightly older DLPack version if we do not want to bump up NumPy/CuPy/PyTorch/... to very recent versions. Do we have access to the 2025 Intel Python tools release on Aurora?

Close #9

Action Items

• start by vibing while preparing dinner, then manually:
• review and finish Array4
• PODVector
• Vector
• ArrayOfStructs
• BaseFab
• SmallMatrix
• SYCL: Implement .to_dpnp / .to_dpctl helper functions
• Update .to_xp functions to use .to_dpnp or .to_dpctl for SYCL GPUs
• Test on CUDA GPU
• Test on ROCm GPU
• Test on SYCL GPU (help wanted)
• Search docs for needed updates.
• Fix DLPack stubs Fix PyCapsule for DLPack pybind/pybind11-stubgen#258 or bind manually in pyAMReX

[roelof-groenewald]

I performed some testing of the new functionality on Perlmutter. After the latest commit, the following appears to work as intended:

def test_mfab_cuda_cupy(mfab_device):
    import cupy as cp

    # AMReX -> cupy
    for mfi in mfab_device:   
        marr_cupy_from_dlpack = cp.from_dlpack(mfab_device.array(mfi))
        marr_cupy_from_dlpack[0, 1, 3, 2] = 5

    for mfi in mfab_device:   
        marr_cupy_from_dlpack = cp.from_dlpack(mfab_device.array(mfi))
        print(marr_cupy_from_dlpack[0, 1, 3, 2])

It executes without failure and prints the modified value 5. Inspection of the DLDevice showed that the device was successfully identified as kDLCUDA. The device id returned 3, which is consistent with Perlmutter's standard rank-to-gpu mapping with just one MPI rank.

[ax3l]

Awesome, then we are nearly there.

Try the dpnp logic for SYCL next?

[roelof-groenewald]

I tested the dlpack functionality on Aurora (SYCL) and it also now produces the expected result. I also modified Array4_to_xp to take into account the GPU backend. We can now successfully access a MultiFab's Array4 from a SYCL device with

for mfi in mfab_device:
    mfab_device.array(mfi).to_dpnp()

[roelof-groenewald]

I compiled WarpX on Aurora using this pyamrex branch. With it I was able to successfully run a multi-GPU simulation that uses fields.py to read MultiFab values 🎉 🚀

[roelof-groenewald]

Just want to flag this since copy=True doesn't yet work in the .to_dpnp() function.

[ax3l]

We need to rebase against development after #455 was merged. I already added the DLDeviceType bindings now and the other PR adds capsule type hints.