Helpers: to_numpy/cupy (#88)

* Helpers: to_numpy/cupy
* Update Stub Files

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/Base/PODVector.cpp

@@ -6,6 +6,7 @@
 #include "pyAMReX.H"
 
 #include <AMReX_PODVector.H>
+#include <AMReX_GpuContainers.H>
 
 #include <sstream>
 
@@ -72,6 +73,16 @@ void make_PODVector(py::module &m, std::string typestr, std::string allocstr)
         .def("resize", py::overload_cast<std::size_t, const T&>(&PODVector_type::resize))
         .def("reserve", &PODVector_type::reserve)
         .def("shrink_to_fit", &PODVector_type::shrink_to_fit)
+        .def("to_host", [](PODVector_type const & pv) {
+            PODVector<T, std::allocator<T>> h_data(pv.size());
+            //py::array_t<T> h_data(pv.size());
+            amrex::Gpu::copy(amrex::Gpu::deviceToHost,
+               pv.begin(), pv.end(),
+               h_data.begin()
+               //h_data.ptr()
+            );
+            return h_data;
+        })
 
         // front
         // back

src/Particle/ArrayOfStructs.cpp

@@ -135,6 +135,18 @@ void make_ArrayOfStructs(py::module &m, std::string allocstr)
         .def("test_sizes", [](){ })
         .def("__setitem__", [](AOSType &aos, int const v,  const ParticleType& p){ aos[v] = p; })
         .def("__getitem__", [](AOSType &aos, int const v){ return aos[v]; }, py::return_value_policy::reference)
+
+        .def("to_host", [](AOSType const & aos) {
+            ArrayOfStructs<T_ParticleType, std::allocator> h_data;
+            h_data.resize(aos.size());
+            //py::array_t<T_ParticleType> h_data(aos.size());
+            amrex::Gpu::copy(amrex::Gpu::deviceToHost,
+               aos.begin(), aos.end(),
+               h_data.begin()
+               //h_data.ptr()
+            );
+            return h_data;
+        })
     ;
 }
 

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/ArrayOfStructs.py

@@ -0,0 +1,91 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+from collections import namedtuple
+
+
+def aos_to_numpy(self, copy=False):
+    """
+    Provide Numpy views into a ArrayOfStructs.
+
+    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 arrays.
+    """
+    import numpy as np
+
+    if self.size() == 0:
+        raise ValueError("AoS is empty.")
+
+    if copy:
+        # This supports a device-to-host copy.
+        #
+        # todo: validate of the to_host() returned object
+        #       lifetime is always managed correctly by
+        #       Python's GC - otherwise copy twice via copy=True
+        return np.array(self.to_host(), copy=False)
+    else:
+        return np.array(self, copy=False)
+
+
+def aos_to_cupy(self, copy=False):
+    """
+    Provide Cupy views into a ArrayOfStructs.
+
+    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 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("AoS is empty.")
+
+    return cp.array(self, copy=copy)
+
+
+def register_AoS_extension(amr):
+    """ArrayOfStructs helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every ArrayOfStructs_* type
+    for _, AoS_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("ArrayOfStructs_"),
+    ):
+        AoS_type.to_numpy = aos_to_numpy
+        AoS_type.to_cupy = aos_to_cupy

src/amrex/MultiFab.py

@@ -0,0 +1,94 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def mf_to_numpy(self, copy=False, order="F"):
+    """
+    Provide a Numpy view into a MultiFab.
+
+    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 np.array
+        A list of numpy n-dimensional arrays, for each local block in the
+        MultiFab.
+    """
+    views = []
+    for mfi in self:
+        views.append(self.array(mfi).to_numpy(copy, order))
+
+    return views
+
+
+def mf_to_cupy(self, copy=False, order="F"):
+    """
+    Provide a Cupy view into a MultiFab.
+
+    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 cupy.array
+        A list of cupy n-dimensional arrays, for each local block in the
+        MultiFab.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    views = []
+    for mfi in self:
+        views.append(self.array(mfi).to_cupy(copy, order))
+
+    return views
+
+
+def register_MultiFab_extension(amr):
+    """MultiFab helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every MultiFab* type
+    for _, MultiFab_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("MultiFab"),
+    ):
+        MultiFab_type.to_numpy = mf_to_numpy
+        MultiFab_type.to_cupy = mf_to_cupy