Helpers: to_numpy/cupy

MANIFEST.in

@@ -6,6 +6,10 @@ recursive-include cmake *
 recursive-include src *
 recursive-include tests *
 
+# avoid accidentially copying compiled Python files
+global-exclude */__pycache__/*
+global-exclude *.pyc
+
 # see .gitignore
 prune cmake-build*
 prune .spack-env*

src/amrex/Array4.py

@@ -0,0 +1,98 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def array4_to_numpy(self, copy=False, order="F"):
+    """
+    Provide a Numpy view into an Array4.
+
+    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
+    -------
+    np.array
+        A numpy n-dimensional array.
+    """
+    import numpy as np
+
+    if order == "F":
+        return np.array(self, copy=copy).T
+    elif order == "C":
+        return np.array(self, copy=copy)
+    else:
+        raise ValueError("The order argument must be F or C.")
+
+
+def array4_to_cupy(self, copy=False, order="F"):
+    """
+    Provide a Cupy view into an Array4.
+
+    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
+    -------
+    cupy.array
+        A cupy n-dimensional array.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    if order == "F":
+        return cp.array(self, copy=copy).T
+    elif order == "C":
+        return cp.array(self, copy=copy)
+    else:
+        raise ValueError("The order argument must be F or C.")
+
+
+def register_Array4_extension(amr):
+    """Array4 helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every Array4_* type
+    for _, Array4_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("Array4_"),
+    ):
+        Array4_type.to_numpy = array4_to_numpy
+        Array4_type.to_cupy = array4_to_cupy

src/amrex/PODVector.py

@@ -0,0 +1,76 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def podvector_to_numpy(self, copy=False):
+    """
+    Provide a Numpy view into a PODVector (e.g., RealVector, IntVector).
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    np.array
+        A 1D numpy array.
+    """
+    import numpy as np
+
+    if self.size() > 0:
+        return np.array(self, copy=copy)
+    else:
+        raise ValueError("Vector is empty.")
+
+
+def podvector_to_cupy(self, copy=False):
+    """
+    Provide a Cupy view into a PODVector (e.g., RealVector, IntVector).
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    cupy.array
+        A 1D cupy array.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    if self.size() > 0:
+        return cp.array(self, copy=copy)
+    else:
+        raise ValueError("Vector is empty.")
+
+
+def register_PODVector_extension(amr):
+    """PODVector helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every PODVector_* type
+    for _, POD_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("PODVector_"),
+    ):
+        POD_type.to_numpy = podvector_to_numpy
+        POD_type.to_cupy = podvector_to_cupy

src/amrex/StructOfArrays.py

@@ -0,0 +1,99 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+from collections import namedtuple
+
+
+def soa_to_numpy(self, copy=False):
+    """
+    Provide Numpy views into a StructOfArrays.
+
+    Parameters
+    ----------
+    self : amrex.StructOfArrays_*
+        A StructOfArrays 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 arrays.
+    """
+    import numpy as np
+
+    SoA_np = namedtuple(type(self).__name__ + "_np", ["real", "int"])
+
+    soa_view = SoA_np([], [])
+
+    if self.size() == 0:
+        raise ValueError("SoA is empty.")
+
+    for idx_real in range(self.NumRealComps()):
+        soa_view.real.append(np.array(self.GetRealData(idx_real), copy=copy))
+
+    for idx_int in range(self.NumIntComps()):
+        soa_view.int.append(np.array(self.GetIntData(idx_int), copy=copy))
+
+    return soa_view
+
+
+def soa_to_cupy(self, copy=False):
+    """
+    Provide Cupy views into a StructOfArrays.
+
+    Parameters
+    ----------
+    self : amrex.StructOfArrays_*
+        A StructOfArrays 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 arrays.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    SoA_cp = namedtuple(type(self).__name__ + "_cp", ["real", "int"])
+
+    soa_view = SoA_cp([], [])
+
+    if self.size() == 0:
+        raise ValueError("SoA is empty.")
+
+    for idx_real in range(self.NumRealComps()):
+        soa_view.real.append(cp.array(self.GetRealData(idx_real), copy=copy))
+
+    for idx_int in range(self.NumIntComps()):
+        soa_view.int.append(cp.array(self.GetIntData(idx_int), copy=copy))
+
+    return soa_view
+
+
+def register_SoA_extension(amr):
+    """StructOfArrays helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every StructOfArrays_* type
+    for _, SoA_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("StructOfArrays_"),
+    ):
+        SoA_type.to_numpy = soa_to_numpy
+        SoA_type.to_cupy = soa_to_cupy

src/amrex/space1d/__init__.py

@@ -41,3 +41,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_1d_pybind)
+register_PODVector_extension(amrex_1d_pybind)
+register_SoA_extension(amrex_1d_pybind)

src/amrex/space2d/__init__.py

@@ -41,3 +41,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_2d_pybind)
+register_PODVector_extension(amrex_2d_pybind)
+register_SoA_extension(amrex_2d_pybind)

src/amrex/space3d/__init__.py

@@ -41,3 +41,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_3d_pybind)
+register_PODVector_extension(amrex_3d_pybind)
+register_SoA_extension(amrex_3d_pybind)

tests/test_multifab.py

@@ -45,26 +45,25 @@ def test_mfab_loop(make_mfab):
 
         # numpy representation: non-copying view, including the
         # guard/ghost region
-        #   note: in numpy, indices are in C-order!
-        marr_np = np.array(marr, copy=False)
+        marr_np = marr.to_numpy()
 
         # check the values at start/end are the same: first component
         assert marr_np[0, 0, 0, 0] == marr[bx.small_end]
-        assert marr_np[0, -1, -1, -1] == marr[bx.big_end]
+        assert marr_np[-1, -1, -1, 0] == marr[bx.big_end]
         # same check, but for all components
         for n in range(mfab.num_comp):
             small_end_comp = list(bx.small_end) + [n]
             big_end_comp = list(bx.big_end) + [n]
-            assert marr_np[n, 0, 0, 0] == marr[small_end_comp]
-            assert marr_np[n, -1, -1, -1] == marr[big_end_comp]
+            assert marr_np[0, 0, 0, n] == marr[small_end_comp]
+            assert marr_np[-1, -1, -1, n] == marr[big_end_comp]
 
         # now we do some faster assignments, using range based access
         #   this should fail as out-of-bounds, but does not
         #     does Numpy not check array access for non-owned views?
         # marr_np[24:200, :, :, :] = 42.
 
         #   all components and all indices set at once to 42
-        marr_np[:, :, :, :] = 42.0
+        marr_np[()] = 42.0
 
         # values in start & end still match?
         assert marr_np[0, 0, 0, 0] == marr[bx.small_end]
@@ -210,10 +209,11 @@ def test_mfab_ops_cuda_cupy(make_mfab_device):
     with cupy.profiler.time_range("assign 3 [()]", color_id=0):
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="C")
             # print(marr_cupy.shape)  # 1, 32, 32, 32
             # print(marr_cupy.dtype)  # float64
+            # performance:
+            #   https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
 
             # write and read into the marr_cupy
             marr_cupy[()] = 3.0
@@ -244,8 +244,11 @@ def set_to_five(mm):
 
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="F")
+            # print(marr_cupy.shape)  # 32, 32, 32, 1
+            # print(marr_cupy.dtype)  # float64
+            # performance:
+            #   https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
 
             # write and read into the marr_cupy
             fives_cp = set_to_five(marr_cupy)
@@ -266,8 +269,7 @@ def set_to_seven(x):
 
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="C")
 
             # write and read into the marr_cupy
             set_to_seven(marr_cupy)