diff --git a/src/amrex/Array4.py b/src/amrex/Array4.py index 36f0ffbe..8cd73f8a 100644 --- a/src/amrex/Array4.py +++ b/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 diff --git a/src/amrex/ArrayOfStructs.py b/src/amrex/ArrayOfStructs.py index 3c2d3791..54afb0d4 100644 --- a/src/amrex/ArrayOfStructs.py +++ b/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 diff --git a/src/amrex/MultiFab.py b/src/amrex/MultiFab.py index 3900626e..ab0f1379 100644 --- a/src/amrex/MultiFab.py +++ b/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__ diff --git a/src/amrex/PODVector.py b/src/amrex/PODVector.py index c241405c..667b52c9 100644 --- a/src/amrex/PODVector.py +++ b/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 diff --git a/src/amrex/StructOfArrays.py b/src/amrex/StructOfArrays.py index 7dff6128..36f0b967 100644 --- a/src/amrex/StructOfArrays.py +++ b/src/amrex/StructOfArrays.py @@ -76,7 +76,7 @@ def soa_int_comps(self, num_comps): def soa_to_numpy(self, copy=False): """ - Provide Numpy views into a StructOfArrays. + Provide NumPy views into a StructOfArrays. Parameters ---------- @@ -89,7 +89,7 @@ def soa_to_numpy(self, copy=False): ------- namedtuple A tuple with real and int components that are each dicts - of 1D numpy arrays. The dictionary key order is the same as + of 1D NumPy arrays. The dictionary key order is the same as in the C++ component order. For pure SoA particle layouts, an additional component idcpu with global particle indices is populated. @@ -129,7 +129,7 @@ def soa_to_numpy(self, copy=False): def soa_to_cupy(self, copy=False): """ - Provide Cupy views into a StructOfArrays. + Provide CuPy views into a StructOfArrays. Parameters ---------- @@ -142,7 +142,7 @@ def soa_to_cupy(self, copy=False): ------- namedtuple A tuple with real and int components that are each dicts - of 1D numpy arrays. The dictionary key order is the same as + of 1D NumPy arrays. The dictionary key order is the same as in the C++ component order. For pure SoA particle layouts, an additional component idcpu with global particle indices is populated. @@ -183,6 +183,35 @@ def soa_to_cupy(self, copy=False): return soa_view +def soa_to_xp(self, copy=False): + """ + Provide NumPy or CuPy views into a StructOfArrays, 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.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 dicts + of 1D NumPy or CuPy arrays. The dictionary key order is the same as + in the C++ component order. + For pure SoA particle layouts, an additional component idcpu + with global particle indices is populated. + """ + import inspect + + amr = inspect.getmodule(self) + return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy) + + def register_SoA_extension(amr): """StructOfArrays helper methods""" import inspect @@ -202,3 +231,4 @@ def register_SoA_extension(amr): # converters SoA_type.to_numpy = soa_to_numpy SoA_type.to_cupy = soa_to_cupy + SoA_type.to_xp = soa_to_xp