to_xp: AoS, SoA, PODVector, MF, Array4 (AMReX-Codes#289)

* `to_xp`: AoS, SoA, PODVector, MF, Array4

Add a `to_xp()` method that returns either NumPy or CuPy
containers, depending on `amr.Config.have_gpu`, to all
major AMReX data containers.

* Update Tests

* Test default: `amrex.the_arena_is_managed=0`

* Documentation: `.to_xp()`

docs/source/usage/zerocopy.rst

@@ -7,7 +7,7 @@ The Python binding pyAMReX bridges the compute in AMReX block-structured codes a
 As such, it includes zero-copy GPU data access for AI/ML, in situ analysis, application coupling by implementing :ref:`standardized data interfaces <developers-implementation>`.
 
 
-CPU: numpy
+CPU: NumPy
 ----------
 
 zero-copy read and write access.
@@ -16,18 +16,30 @@ 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.
+Writing to the created NumPy array will also modify the underlying AMReX memory.
 
 
-GPU: cupy
+GPU: CuPy
 ---------
 
 GPU zero-copy read and write access.
 
 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.
+Writing to the created CuPy array will also modify the underlying AMReX memory.
+
+
+CPU/GPU Agnostic Code: NumPy/CuPy
+---------------------------------
+
+The previous examples can be written in CPU/GPU agnostics manner.
+Either using NumPy (``np``) or CuPy (``cp``), we provide a `common short-hand abbreviation <https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code>`__ named ``xp`` .
+
+Call ``.to_xp()`` on data objects of pyAMReX.
+See the optional arguments of this API.
+
+Writing to the created NumPy/CuPy array will also modify the underlying AMReX memory.
 
 
 GPU: numba

src/amrex/Array4.py

@@ -9,7 +9,7 @@
 
 def array4_to_numpy(self, copy=False, order="F"):
     """
-    Provide a Numpy view into an Array4.
+    Provide a NumPy view into an Array4.
 
     This includes ngrow guard cells of the box.
 
@@ -32,7 +32,7 @@ def array4_to_numpy(self, copy=False, order="F"):
     Returns
     -------
     np.array
-        A numpy n-dimensional array.
+        A NumPy n-dimensional array.
     """
     import numpy as np
 
@@ -52,7 +52,7 @@ def array4_to_numpy(self, copy=False, order="F"):
 
 def array4_to_cupy(self, copy=False, order="F"):
     """
-    Provide a Cupy view into an Array4.
+    Provide a CuPy view into an Array4.
 
     This includes ngrow guard cells of the box.
 
@@ -92,6 +92,44 @@ def array4_to_cupy(self, copy=False, order="F"):
         raise ValueError("The order argument must be F or C.")
 
 
+def array4_to_xp(self, copy=False, order="F"):
+    """
+    Provide a NumPy or CuPy view into an Array4, depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    This includes ngrow guard cells of the box.
+
+    Note on the order of indices:
+    By default, this is as in AMReX in Fortran contiguous order, indexing as
+    x,y,z. This has performance implications for use in external libraries such
+    as cupy.
+    The order="C" option will index as z,y,x and perform better with cupy.
+    https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
+
+    Parameters
+    ----------
+    self : amrex.Array4_*
+        An Array4 class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+    order : string, optional
+        F order (default) or C. C is faster with external libraries.
+
+    Returns
+    -------
+    xp.array
+        A NumPy or CuPy n-dimensional array.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return (
+        self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order)
+    )
+
+
 def register_Array4_extension(amr):
     """Array4 helper methods"""
     import inspect
@@ -106,3 +144,4 @@ def register_Array4_extension(amr):
     ):
         Array4_type.to_numpy = array4_to_numpy
         Array4_type.to_cupy = array4_to_cupy
+        Array4_type.to_xp = array4_to_xp

src/amrex/ArrayOfStructs.py

@@ -9,7 +9,7 @@
 
 def aos_to_numpy(self, copy=False):
     """
-    Provide Numpy views into a ArrayOfStructs.
+    Provide NumPy views into a ArrayOfStructs.
 
     Parameters
     ----------
@@ -22,7 +22,7 @@ def aos_to_numpy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each lists
-        of 1D numpy arrays.
+        of 1D NumPy arrays.
     """
     import numpy as np
 
@@ -42,7 +42,7 @@ def aos_to_numpy(self, copy=False):
 
 def aos_to_cupy(self, copy=False):
     """
-    Provide Cupy views into a ArrayOfStructs.
+    Provide CuPy views into a ArrayOfStructs.
 
     Parameters
     ----------
@@ -55,7 +55,7 @@ def aos_to_cupy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each lists
-        of 1D numpy arrays.
+        of 1D NumPy arrays.
 
     Raises
     ------
@@ -70,6 +70,32 @@ def aos_to_cupy(self, copy=False):
     return cp.array(self, copy=copy)
 
 
+def aos_to_xp(self, copy=False):
+    """
+    Provide NumPy or CuPy views into a ArrayOfStructs, depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    self : amrex.ArrayOfStructs_*
+        An ArrayOfStructs class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    namedtuple
+        A tuple with real and int components that are each lists
+        of 1D NumPy or CuPy arrays.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
+
+
 def register_AoS_extension(amr):
     """ArrayOfStructs helper methods"""
     import inspect
@@ -84,3 +110,4 @@ def register_AoS_extension(amr):
     ):
         AoS_type.to_numpy = aos_to_numpy
         AoS_type.to_cupy = aos_to_cupy
+        AoS_type.to_xp = aos_to_xp

src/amrex/MultiFab.py

@@ -9,9 +9,9 @@
 from .Iterator import next
 
 
-def mf_to_numpy(amr, self, copy=False, order="F"):
+def mf_to_numpy(self, copy=False, order="F"):
     """
-    Provide a Numpy view into a MultiFab.
+    Provide a NumPy view into a MultiFab.
 
     This includes ngrow guard cells of each box.
 
@@ -34,9 +34,13 @@ def mf_to_numpy(amr, self, copy=False, order="F"):
     Returns
     -------
     list of numpy.array
-        A list of numpy n-dimensional arrays, for each local block in the
+        A list of NumPy n-dimensional arrays, for each local block in the
         MultiFab.
     """
+    import inspect
+
+    amr = inspect.getmodule(self)
+
     mf = self
     if copy:
         mf = amr.MultiFab(
@@ -58,7 +62,7 @@ def mf_to_numpy(amr, self, copy=False, order="F"):
 
 def mf_to_cupy(self, copy=False, order="F"):
     """
-    Provide a Cupy view into a MultiFab.
+    Provide a CuPy view into a MultiFab.
 
     This includes ngrow guard cells of each box.
 
@@ -81,7 +85,7 @@ def mf_to_cupy(self, copy=False, order="F"):
     Returns
     -------
     list of cupy.array
-        A list of cupy n-dimensional arrays, for each local block in the
+        A list of CuPy n-dimensional arrays, for each local block in the
         MultiFab.
 
     Raises
@@ -96,6 +100,46 @@ def mf_to_cupy(self, copy=False, order="F"):
     return views
 
 
+def mf_to_xp(self, copy=False, order="F"):
+    """
+    Provide a NumPy or CuPy view into a MultiFab,
+    depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    This includes ngrow guard cells of each box.
+
+    Note on the order of indices:
+    By default, this is as in AMReX in Fortran contiguous order, indexing as
+    x,y,z. This has performance implications for use in external libraries such
+    as cupy.
+    The order="C" option will index as z,y,x and perform better with cupy.
+    https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
+
+    Parameters
+    ----------
+    self : amrex.MultiFab
+        A MultiFab class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+    order : string, optional
+        F order (default) or C. C is faster with external libraries.
+
+    Returns
+    -------
+    list of xp.array
+        A list of NumPy or CuPy n-dimensional arrays, for each local block in the
+        MultiFab.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return (
+        self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order)
+    )
+
+
 def copy_multifab(amr, self):
     """
     Create a copy of this MultiFab, using the same Arena.
@@ -141,12 +185,9 @@ def register_MultiFab_extension(amr):
     # register member functions for the MultiFab type
     amr.MultiFab.__iter__ = lambda mfab: amr.MFIter(mfab)
 
-    amr.MultiFab.to_numpy = lambda self, copy=False, order="F": mf_to_numpy(
-        amr, self, copy, order
-    )
-    amr.MultiFab.to_numpy.__doc__ = mf_to_numpy.__doc__
-
+    amr.MultiFab.to_numpy = mf_to_numpy
     amr.MultiFab.to_cupy = mf_to_cupy
+    amr.MultiFab.to_xp = mf_to_xp
 
     amr.MultiFab.copy = lambda self: copy_multifab(amr, self)
     amr.MultiFab.copy.__doc__ = copy_multifab.__doc__

src/amrex/PODVector.py

@@ -9,7 +9,7 @@
 
 def podvector_to_numpy(self, copy=False):
     """
-    Provide a Numpy view into a PODVector (e.g., RealVector, IntVector).
+    Provide a NumPy view into a PODVector (e.g., RealVector, IntVector).
 
     Parameters
     ----------
@@ -21,7 +21,7 @@ def podvector_to_numpy(self, copy=False):
     Returns
     -------
     np.array
-        A 1D numpy array.
+        A 1D NumPy array.
     """
     import numpy as np
 
@@ -41,7 +41,7 @@ def podvector_to_numpy(self, copy=False):
 
 def podvector_to_cupy(self, copy=False):
     """
-    Provide a Cupy view into a PODVector (e.g., RealVector, IntVector).
+    Provide a CuPy view into a PODVector (e.g., RealVector, IntVector).
 
     Parameters
     ----------
@@ -68,6 +68,32 @@ def podvector_to_cupy(self, copy=False):
         raise ValueError("Vector is empty.")
 
 
+def podvector_to_xp(self, copy=False):
+    """
+    Provide a NumPy or CuPy view into a PODVector (e.g., RealVector, IntVector),
+    depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    xp.array
+        A 1D NumPy or CuPy array.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
+
+
 def register_PODVector_extension(amr):
     """PODVector helper methods"""
     import inspect
@@ -82,3 +108,4 @@ def register_PODVector_extension(amr):
     ):
         POD_type.to_numpy = podvector_to_numpy
         POD_type.to_cupy = podvector_to_cupy
+        POD_type.to_xp = podvector_to_xp