From 9aa7cb88b01f9b8d18f7aa88bccfae2c6360c107 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 22 Sep 2023 22:40:26 +0200 Subject: [PATCH] MultiFab: to_numpy/cupy (#192) * MultiFab: to_numpy/cupy Add numpy & cupy helpers for MultiFab. * Update Stub Files --------- Co-authored-by: ax3l --- src/amrex/MultiFab.py | 39 +++++++----- src/amrex/space1d/__init__.py | 2 + src/amrex/space1d/__init__.pyi | 2 + .../space1d/amrex_1d_pybind/__init__.pyi | 61 +++++++++++++++++++ src/amrex/space2d/__init__.py | 2 + src/amrex/space2d/__init__.pyi | 2 + .../space2d/amrex_2d_pybind/__init__.pyi | 61 +++++++++++++++++++ src/amrex/space3d/__init__.py | 2 + src/amrex/space3d/__init__.pyi | 2 + .../space3d/amrex_3d_pybind/__init__.pyi | 61 +++++++++++++++++++ tests/test_multifab.py | 10 +++ 11 files changed, 228 insertions(+), 16 deletions(-) diff --git a/src/amrex/MultiFab.py b/src/amrex/MultiFab.py index 30ca5f4a..a2b46899 100644 --- a/src/amrex/MultiFab.py +++ b/src/amrex/MultiFab.py @@ -7,7 +7,7 @@ """ -def mf_to_numpy(self, copy=False, order="F"): +def mf_to_numpy(amr, self, copy=False, order="F"): """ Provide a Numpy view into a MultiFab. @@ -29,13 +29,24 @@ def mf_to_numpy(self, copy=False, order="F"): Returns ------- - list of np.array + list of numpy.array A list of numpy n-dimensional arrays, for each local block in the MultiFab. """ + mf = self + if copy: + mf = amr.MultiFab( + self.box_array(), + self.dm(), + self.n_comp(), + self.n_grow_vect(), + amr.MFInfo().set_arena(amr.The_Pinned_Arena()), + ) + amr.dtoh_memcpy(mf, self) + views = [] - for mfi in self: - views.append(self.array(mfi).to_numpy(copy, order)) + for mfi in mf: + views.append(mf.array(mfi).to_numpy(copy=False, order=order)) return views @@ -80,15 +91,11 @@ def mf_to_cupy(self, copy=False, order="F"): 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 + + # register member functions for the MultiFab type + 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_cupy = mf_to_cupy diff --git a/src/amrex/space1d/__init__.py b/src/amrex/space1d/__init__.py index 76b75647..d752c667 100644 --- a/src/amrex/space1d/__init__.py +++ b/src/amrex/space1d/__init__.py @@ -46,10 +46,12 @@ def Print(*args, **kwargs): from ..Array4 import register_Array4_extension from ..ArrayOfStructs import register_AoS_extension +from ..MultiFab import register_MultiFab_extension from ..PODVector import register_PODVector_extension from ..StructOfArrays import register_SoA_extension register_Array4_extension(amrex_1d_pybind) +register_MultiFab_extension(amrex_1d_pybind) register_PODVector_extension(amrex_1d_pybind) register_SoA_extension(amrex_1d_pybind) register_AoS_extension(amrex_1d_pybind) diff --git a/src/amrex/space1d/__init__.pyi b/src/amrex/space1d/__init__.pyi index d59d128f..e435534b 100644 --- a/src/amrex/space1d/__init__.pyi +++ b/src/amrex/space1d/__init__.pyi @@ -38,6 +38,7 @@ import os as os from amrex.Array4 import register_Array4_extension from amrex.ArrayOfStructs import register_AoS_extension +from amrex.MultiFab import register_MultiFab_extension from amrex.PODVector import register_PODVector_extension from amrex.StructOfArrays import register_SoA_extension from amrex.space1d.amrex_1d_pybind import ( @@ -461,6 +462,7 @@ __all__ = [ "refine", "register_AoS_extension", "register_Array4_extension", + "register_MultiFab_extension", "register_PODVector_extension", "register_SoA_extension", "size", diff --git a/src/amrex/space1d/amrex_1d_pybind/__init__.pyi b/src/amrex/space1d/amrex_1d_pybind/__init__.pyi index d10686c6..437efff5 100644 --- a/src/amrex/space1d/amrex_1d_pybind/__init__.pyi +++ b/src/amrex/space1d/amrex_1d_pybind/__init__.pyi @@ -4043,6 +4043,67 @@ class MultiFab(FabArray_FArrayBox): """ Same as sum with local=false, but for non-cell-centered data, thisskips non-unique points that are owned by multiple boxes. """ + def 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 + + """ + def 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 numpy.array + A list of numpy n-dimensional arrays, for each local block in the + MultiFab. + + """ def weighted_sync(self, arg0: MultiFab, arg1: Periodicity) -> None: ... class PIdx: diff --git a/src/amrex/space2d/__init__.py b/src/amrex/space2d/__init__.py index b3626276..d0933662 100644 --- a/src/amrex/space2d/__init__.py +++ b/src/amrex/space2d/__init__.py @@ -46,10 +46,12 @@ def Print(*args, **kwargs): from ..Array4 import register_Array4_extension from ..ArrayOfStructs import register_AoS_extension +from ..MultiFab import register_MultiFab_extension from ..PODVector import register_PODVector_extension from ..StructOfArrays import register_SoA_extension register_Array4_extension(amrex_2d_pybind) +register_MultiFab_extension(amrex_2d_pybind) register_PODVector_extension(amrex_2d_pybind) register_SoA_extension(amrex_2d_pybind) register_AoS_extension(amrex_2d_pybind) diff --git a/src/amrex/space2d/__init__.pyi b/src/amrex/space2d/__init__.pyi index 9715ccd0..b293fc6d 100644 --- a/src/amrex/space2d/__init__.pyi +++ b/src/amrex/space2d/__init__.pyi @@ -38,6 +38,7 @@ import os as os from amrex.Array4 import register_Array4_extension from amrex.ArrayOfStructs import register_AoS_extension +from amrex.MultiFab import register_MultiFab_extension from amrex.PODVector import register_PODVector_extension from amrex.StructOfArrays import register_SoA_extension from amrex.space2d.amrex_2d_pybind import ( @@ -461,6 +462,7 @@ __all__ = [ "refine", "register_AoS_extension", "register_Array4_extension", + "register_MultiFab_extension", "register_PODVector_extension", "register_SoA_extension", "size", diff --git a/src/amrex/space2d/amrex_2d_pybind/__init__.pyi b/src/amrex/space2d/amrex_2d_pybind/__init__.pyi index 45d0863d..56d84fac 100644 --- a/src/amrex/space2d/amrex_2d_pybind/__init__.pyi +++ b/src/amrex/space2d/amrex_2d_pybind/__init__.pyi @@ -4049,6 +4049,67 @@ class MultiFab(FabArray_FArrayBox): """ Same as sum with local=false, but for non-cell-centered data, thisskips non-unique points that are owned by multiple boxes. """ + def 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 + + """ + def 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 numpy.array + A list of numpy n-dimensional arrays, for each local block in the + MultiFab. + + """ def weighted_sync(self, arg0: MultiFab, arg1: Periodicity) -> None: ... class PIdx: diff --git a/src/amrex/space3d/__init__.py b/src/amrex/space3d/__init__.py index 4b4fd623..269b8112 100644 --- a/src/amrex/space3d/__init__.py +++ b/src/amrex/space3d/__init__.py @@ -46,10 +46,12 @@ def Print(*args, **kwargs): from ..Array4 import register_Array4_extension from ..ArrayOfStructs import register_AoS_extension +from ..MultiFab import register_MultiFab_extension from ..PODVector import register_PODVector_extension from ..StructOfArrays import register_SoA_extension register_Array4_extension(amrex_3d_pybind) +register_MultiFab_extension(amrex_3d_pybind) register_PODVector_extension(amrex_3d_pybind) register_SoA_extension(amrex_3d_pybind) register_AoS_extension(amrex_3d_pybind) diff --git a/src/amrex/space3d/__init__.pyi b/src/amrex/space3d/__init__.pyi index 6a68ed0a..85ec9a12 100644 --- a/src/amrex/space3d/__init__.pyi +++ b/src/amrex/space3d/__init__.pyi @@ -38,6 +38,7 @@ import os as os from amrex.Array4 import register_Array4_extension from amrex.ArrayOfStructs import register_AoS_extension +from amrex.MultiFab import register_MultiFab_extension from amrex.PODVector import register_PODVector_extension from amrex.StructOfArrays import register_SoA_extension from amrex.space3d.amrex_3d_pybind import ( @@ -461,6 +462,7 @@ __all__ = [ "refine", "register_AoS_extension", "register_Array4_extension", + "register_MultiFab_extension", "register_PODVector_extension", "register_SoA_extension", "size", diff --git a/src/amrex/space3d/amrex_3d_pybind/__init__.pyi b/src/amrex/space3d/amrex_3d_pybind/__init__.pyi index f5a30e7d..ebf87c61 100644 --- a/src/amrex/space3d/amrex_3d_pybind/__init__.pyi +++ b/src/amrex/space3d/amrex_3d_pybind/__init__.pyi @@ -4052,6 +4052,67 @@ class MultiFab(FabArray_FArrayBox): """ Same as sum with local=false, but for non-cell-centered data, thisskips non-unique points that are owned by multiple boxes. """ + def 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 + + """ + def 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 numpy.array + A list of numpy n-dimensional arrays, for each local block in the + MultiFab. + + """ def weighted_sync(self, arg0: MultiFab, arg1: Periodicity) -> None: ... class PIdx: diff --git a/tests/test_multifab.py b/tests/test_multifab.py index 80727615..595fef85 100644 --- a/tests/test_multifab.py +++ b/tests/test_multifab.py @@ -350,3 +350,13 @@ def test_mfab_dtoh_copy(make_mfab_device): device_max = mfab_device.max(0) assert device_min == device_max assert device_max == 11.0 + + # numpy bindings (w/ copy) + local_boxes_host = mfab_device.to_numpy(copy=True) + assert max([np.max(box) for box in local_boxes_host]) == device_max + + # cupy bindings (w/o copy) + import cupy as cp + + local_boxes_device = mfab_device.to_cupy() + assert max([cp.max(box) for box in local_boxes_device]) == device_max