From d8da7df53e3543201a70bed7ce5ed20ea80a614c Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Thu, 20 Oct 2022 21:39:22 -0700 Subject: [PATCH] Helpers: to_numpy/cupy --- MANIFEST.in | 4 ++ src/amrex/Array4.py | 98 ++++++++++++++++++++++++++++++++++ src/amrex/PODVector.py | 76 +++++++++++++++++++++++++++ src/amrex/StructOfArrays.py | 99 +++++++++++++++++++++++++++++++++++ src/amrex/space1d/__init__.py | 10 ++++ src/amrex/space2d/__init__.py | 10 ++++ src/amrex/space3d/__init__.py | 10 ++++ tests/test_multifab.py | 26 ++++----- 8 files changed, 321 insertions(+), 12 deletions(-) create mode 100644 src/amrex/Array4.py create mode 100644 src/amrex/PODVector.py create mode 100644 src/amrex/StructOfArrays.py diff --git a/MANIFEST.in b/MANIFEST.in index 0b0abb03..9a61a6d7 100644 --- a/MANIFEST.in +++ b/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* diff --git a/src/amrex/Array4.py b/src/amrex/Array4.py new file mode 100644 index 00000000..5bc28448 --- /dev/null +++ b/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 diff --git a/src/amrex/PODVector.py b/src/amrex/PODVector.py new file mode 100644 index 00000000..ac807182 --- /dev/null +++ b/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 diff --git a/src/amrex/StructOfArrays.py b/src/amrex/StructOfArrays.py new file mode 100644 index 00000000..3926932c --- /dev/null +++ b/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 diff --git a/src/amrex/space1d/__init__.py b/src/amrex/space1d/__init__.py index 6e042339..ad33bec2 100644 --- a/src/amrex/space1d/__init__.py +++ b/src/amrex/space1d/__init__.py @@ -31,6 +31,7 @@ # in pure Python or add some other Python logic # def d_decl(x, y, z): + """Return a tuple of the first passed element""" return (x,) @@ -41,3 +42,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) diff --git a/src/amrex/space2d/__init__.py b/src/amrex/space2d/__init__.py index 235189dd..6718be7b 100644 --- a/src/amrex/space2d/__init__.py +++ b/src/amrex/space2d/__init__.py @@ -31,6 +31,7 @@ # in pure Python or add some other Python logic # def d_decl(x, y, z): + """Return a tuple of the first two passed elements""" return (x, y) @@ -41,3 +42,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) diff --git a/src/amrex/space3d/__init__.py b/src/amrex/space3d/__init__.py index 50e7897c..80880f2a 100644 --- a/src/amrex/space3d/__init__.py +++ b/src/amrex/space3d/__init__.py @@ -31,6 +31,7 @@ # in pure Python or add some other Python logic # def d_decl(x, y, z): + """Return a tuple of the three passed elements""" return (x, y, z) @@ -41,3 +42,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) diff --git a/tests/test_multifab.py b/tests/test_multifab.py index 7ad60c0d..80727615 100644 --- a/tests/test_multifab.py +++ b/tests/test_multifab.py @@ -45,18 +45,17 @@ 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 @@ -64,7 +63,7 @@ def test_mfab_loop(make_mfab): # 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)