AMReX-Codes / pyamrex

GPU-Enabled, Zero-Copy AMReX Python Bindings including AI/ML
http://pyamrex.readthedocs.io
Other
32 stars 15 forks source link

Helper for numpy/cupy based on `amr.Config` #216

Closed ax3l closed 2 months ago

ax3l commented 8 months ago

Let's add a helper for this pattern:

field_list = mfab.to_cupy() if Config.have_gpu else mfab.to_numpy()

for pti in pc.iterator(pc, level=0):
    soa = pti.soa().to_cupy() if Config.have_gpu else pti.soa().to_numpy()

or this https://github.com/AMReX-Codes/amrex-tutorials/blob/58fff188ab8963702db208ac6ff0bd47188ceeaa/GuidedTutorials/HeatEquation/Source/main.py#L14-L33

To simplify writing host-device agnostic Python code like this: https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code

Maybe we can add

field_list = mfab.to_xp()

for pti in pc.iterator(pc, level=0):
    soa = pti.soa().to_xp()
ax3l commented 3 months ago

Feature will land via #289