Helpers: to_numpy/cupy

Author: ax3l

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,88 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2022 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 numpy 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(Array4_type):
+    """Array4 helper methods"""
+    Array4_type.to_numpy = array4_to_numpy
+    Array4_type.to_cupy = array4_to_cupy

src/amrex/__init__.py

@@ -1,7 +1,5 @@
 # -*- coding: utf-8 -*-
 
 # __version__ is TODO - only in spaceNd submodules
-__author__ = (
-    "Axel Huebl, Ryan Sandberg, Shreyas Ananthan, Remi Lehe, " "Weiqun Zhang, et al."
-)
+__author__ = "Axel Huebl, Ryan Sandberg, Shreyas Ananthan, Remi Lehe, Andrew Myers, Reva Jambunathan, Edodardo Zoni, Weiqun Zhang"
 __license__ = "BSD-3-Clause-LBNL"

src/amrex/space1d/__init__.py

@@ -32,3 +32,20 @@
 #
 def d_decl(x, y, z):
     return (x,)
+
+
+from ..Array4 import register_Array4_extension
+
+register_Array4_extension(Array4_float)
+register_Array4_extension(Array4_double)
+register_Array4_extension(Array4_longdouble)
+
+register_Array4_extension(Array4_short)
+register_Array4_extension(Array4_int)
+register_Array4_extension(Array4_long)
+register_Array4_extension(Array4_longlong)
+
+register_Array4_extension(Array4_ushort)
+register_Array4_extension(Array4_uint)
+register_Array4_extension(Array4_ulong)
+register_Array4_extension(Array4_ulonglong)

src/amrex/space2d/__init__.py

@@ -32,3 +32,20 @@
 #
 def d_decl(x, y, z):
     return (x, y)
+
+
+from ..Array4 import register_Array4_extension
+
+register_Array4_extension(Array4_float)
+register_Array4_extension(Array4_double)
+register_Array4_extension(Array4_longdouble)
+
+register_Array4_extension(Array4_short)
+register_Array4_extension(Array4_int)
+register_Array4_extension(Array4_long)
+register_Array4_extension(Array4_longlong)
+
+register_Array4_extension(Array4_ushort)
+register_Array4_extension(Array4_uint)
+register_Array4_extension(Array4_ulong)
+register_Array4_extension(Array4_ulonglong)

src/amrex/space3d/__init__.py

@@ -32,3 +32,20 @@
 #
 def d_decl(x, y, z):
     return (x, y, z)
+
+
+from ..Array4 import register_Array4_extension
+
+register_Array4_extension(Array4_float)
+register_Array4_extension(Array4_double)
+register_Array4_extension(Array4_longdouble)
+
+register_Array4_extension(Array4_short)
+register_Array4_extension(Array4_int)
+register_Array4_extension(Array4_long)
+register_Array4_extension(Array4_longlong)
+
+register_Array4_extension(Array4_ushort)
+register_Array4_extension(Array4_uint)
+register_Array4_extension(Array4_ulong)
+register_Array4_extension(Array4_ulonglong)

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)