Array4: __cuda_array_interface__ v3 (#30)

* Array4: __cuda_array_interface__ v2

Start implementing the `__cuda_array_interface__` for zero-copy
data exchange on Nvidia CUDA GPUs.

* MultiFab: CuPy Test

* `MFIter`: `Finalize()` on `StopIteration`

Since `for` loops create no scope in Python, we need to trigger
finalize logic, including stream syncs, before the destructor of
`MultiFab` iterators are called.

* Add numba test

incl. 3D kernel launch

* Add pytorch

* CuPy Fuse: Avoid Extra Memset

* MultiFab Device Test: Fixes

* Update to v3

* Array4: TODO from CUDA

A bit tricky to implement this caster as new constructor.
Not currently needed, but adds comments where to do this.

Co-authored-by: Remi Lehe <[email protected]>

Author: ax3l

src/Base/Array4.cpp

@@ -19,6 +19,61 @@ namespace py = pybind11;
 using namespace amrex;
 
 
+namespace
+{
+    /** CPU: __array_interface__ v3
+     *
+     * https://numpy.org/doc/stable/reference/arrays.interface.html
+     */
+    template<typename T>
+    py::dict
+    array_interface(Array4<T> const & a4)
+    {
+        auto d = py::dict();
+        auto const len = length(a4);
+        // F->C index conversion here
+        // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride];
+        // Buffer dimensions: zero-size shall not skip dimension
+        auto shape = py::make_tuple(
+                a4.ncomp,
+                len.z <= 0 ? 1 : len.z,
+                len.y <= 0 ? 1 : len.y,
+                len.x <= 0 ? 1 : len.x  // fastest varying index
+        );
+        // buffer protocol strides are in bytes, AMReX strides are elements
+        auto const strides = py::make_tuple(
+                sizeof(T) * a4.nstride,
+                sizeof(T) * a4.kstride,
+                sizeof(T) * a4.jstride,
+                sizeof(T)  // fastest varying index
+        );
+        bool const read_only = false;
+        d["data"] = py::make_tuple(std::intptr_t(a4.dataPtr()), read_only);
+        // note: if we want to keep the same global indexing with non-zero
+        //       box small_end as in AMReX, then we can explore playing with
+        //       this offset as well
+        //d["offset"] = 0;         // default
+        //d["mask"] = py::none();  // default
+
+        d["shape"] = shape;
+        // we could also set this after checking the strides are C-style contiguous:
+        //if (is_contiguous<T>(shape, strides))
+        //    d["strides"] = py::none();  // C-style contiguous
+        //else
+        d["strides"] = strides;
+
+        // type description
+        // for more complicated types, e.g., tuples/structs
+        //d["descr"] = ...;
+        // we currently only need this
+        d["typestr"] = py::format_descriptor<T>::format();
+
+        d["version"] = 3;
+        return d;
+    }
+}
+
+
 template< typename T >
 void make_Array4(py::module &m, std::string typestr)
 {
@@ -85,56 +140,55 @@ void make_Array4(py::module &m, std::string typestr)
             return a4;
         }))
 
+        /* init from __cuda_array_interface__: non-owning view
+         * TODO
+         */
+
+
+        // CPU: __array_interface__ v3
+        // https://numpy.org/doc/stable/reference/arrays.interface.html
         .def_property_readonly("__array_interface__", [](Array4<T> const & a4) {
-            auto d = py::dict();
-            auto const len = length(a4);
-            // F->C index conversion here
-            // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride];
-            // Buffer dimensions: zero-size shall not skip dimension
-            auto shape = py::make_tuple(
-                    a4.ncomp,
-                    len.z <= 0 ? 1 : len.z,
-                    len.y <= 0 ? 1 : len.y,
-                    len.x <= 0 ? 1 : len.x  // fastest varying index
-            );
-            // buffer protocol strides are in bytes, AMReX strides are elements
-            auto const strides = py::make_tuple(
-                    sizeof(T) * a4.nstride,
-                    sizeof(T) * a4.kstride,
-                    sizeof(T) * a4.jstride,
-                    sizeof(T)  // fastest varying index
-            );
-            bool const read_only = false;
-            d["data"] = py::make_tuple(std::intptr_t(a4.dataPtr()), read_only);
-            // note: if we want to keep the same global indexing with non-zero
-            //       box small_end as in AMReX, then we can explore playing with
-            //       this offset as well
-            //d["offset"] = 0;         // default
-            //d["mask"] = py::none();  // default
-
-            d["shape"] = shape;
-            // we could also set this after checking the strides are C-style contiguous:
-            //if (is_contiguous<T>(shape, strides))
-            //    d["strides"] = py::none();  // C-style contiguous
-            //else
-            d["strides"] = strides;
-
-            d["typestr"] = py::format_descriptor<T>::format();
-            d["version"] = 3;
-            return d;
+            return array_interface(a4);
         })
 
+        // CPU: __array_function__ interface (TODO)
+        //
+        // NEP 18 — A dispatch mechanism for NumPy's high level array functions.
+        //   https://numpy.org/neps/nep-0018-array-function-protocol.html
+        // This enables code using NumPy to be directly operated on Array4 arrays.
+        // __array_function__ feature requires NumPy 1.16 or later.
+
 
-        // TODO: __cuda_array_interface__
+        // Nvidia GPUs: __cuda_array_interface__ v2
         // https://numba.readthedocs.io/en/latest/cuda/cuda_array_interface.html
+        .def_property_readonly("__cuda_array_interface__", [](Array4<T> const & a4) {
+            auto d = array_interface(a4);
+
+            // data:
+            // Because the user of the interface may or may not be in the same context, the most common case is to use cuPointerGetAttribute with CU_POINTER_ATTRIBUTE_DEVICE_POINTER in the CUDA driver API (or the equivalent CUDA Runtime API) to retrieve a device pointer that is usable in the currently active context.
+            // TODO For zero-size arrays, use 0 here.
+
+            // None or integer
+            // An optional stream upon which synchronization must take place at the point of consumption, either by synchronizing on the stream or enqueuing operations on the data on the given stream. Integer values in this entry are as follows:
+            //   0: This is disallowed as it would be ambiguous between None and the default stream, and also between the legacy and per-thread default streams. Any use case where 0 might be given should either use None, 1, or 2 instead for clarity.
+            //   1: The legacy default stream.
+            //   2: The per-thread default stream.
+            //   Any other integer: a cudaStream_t represented as a Python integer.
+            //   When None, no synchronization is required.
+            d["stream"] = py::none();
+
+            d["version"] = 3;
+            return d;
+        })
 
 
-        // TODO: __dlpack__
+        // TODO: __dlpack__ __dlpack_device__
         // DLPack protocol (CPU, NVIDIA GPU, AMD GPU, Intel GPU, etc.)
         // https://dmlc.github.io/dlpack/latest/
         // https://data-apis.org/array-api/latest/design_topics/data_interchange.html
         // https://github.com/data-apis/consortium-feedback/issues/1
         // https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h
+        // https://docs.cupy.dev/en/stable/user_guide/interoperability.html#dlpack-data-exchange-protocol
 
 
         .def("contains", &Array4<T>::contains)

src/Base/MultiFab.cpp

@@ -94,6 +94,7 @@ void init_MultiFab(py::module &m) {
                 if( !mfi.isValid() )
                 {
                     first_or_done = true;
+                    mfi.Finalize();
                     throw py::stop_iteration();
                 }
                 return mfi;

tests/conftest.py

@@ -84,3 +84,31 @@ def create():
         return mfab
 
     return create
+
+
+@pytest.mark.skipif(
+    amrex.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA"
+)
+@pytest.fixture(scope="function", params=list(itertools.product([1, 3], [0, 1])))
+def make_mfab_device(boxarr, distmap, request):
+    """MultiFab that resides purely on the device:
+    The MultiFab object itself is not a fixture because we want to avoid caching
+    it between amrex.initialize/finalize calls of various tests.
+    https://github.com/pytest-dev/pytest/discussions/10387
+    https://github.com/pytest-dev/pytest/issues/5642#issuecomment-1279612764
+    """
+
+    def create():
+        num_components = request.param[0]
+        num_ghost = request.param[1]
+        mfab = amrex.MultiFab(
+            boxarr,
+            distmap,
+            num_components,
+            num_ghost,
+            amrex.MFInfo().set_arena(amrex.The_Device_Arena()),
+        )
+        mfab.set_val(0.0, 0, num_components)
+        return mfab
+
+    return create

tests/test_array4.py

@@ -69,36 +69,81 @@ def test_array4():
     x[1, 1, 1] = 44
     assert v_carr2np[0, 1, 1, 1] == 44
 
-    # from cupy
-
-    # to numpy
-
-    # to cupy
-
-    return
-
-    # Check indexing
-    assert obj[0] == 1
-    assert obj[1] == 2
-    assert obj[2] == 3
-    assert obj[-1] == 3
-    assert obj[-2] == 2
-    assert obj[-3] == 1
-    with pytest.raises(IndexError):
-        obj[-4]
-    with pytest.raises(IndexError):
-        obj[3]
-
-    # Check assignment
-    obj[0] = 2
-    obj[1] = 3
-    obj[2] = 4
-    assert obj[0] == 2
-    assert obj[1] == 3
-    assert obj[2] == 4
-
-
-# def test_iv_conversions():
-#    obj = amrex.IntVect.max_vector().numpy()
-#    assert(isinstance(obj, np.ndarray))
-#    assert(obj.dtype == np.int32)
+
+@pytest.mark.skipif(
+    amrex.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA"
+)
+def test_array4_numba():
+    # https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html
+    from numba import cuda
+
+    # numba -> AMReX Array4
+    x = np.ones(
+        (
+            2,
+            3,
+            4,
+        )
+    )  # type: numpy.ndarray
+
+    # host-to-device copy
+    x_numba = cuda.to_device(x)  # type: numba.cuda.cudadrv.devicearray.DeviceNDArray
+    # x_cupy = cupy.asarray(x_numba)      # type: cupy.ndarray
+    x_arr = amrex.Array4_double(x_numba)  # type: amrex.Array4_double
+
+    assert (
+        x_arr.__cuda_array_interface__["data"][0]
+        == x_numba.__cuda_array_interface__["data"][0]
+    )
+
+    # AMReX -> numba
+    # arr_numba = cuda.as_cuda_array(arr4)
+    # ... or as MultiFab test
+    # TODO
+
+
+@pytest.mark.skipif(
+    amrex.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA"
+)
+def test_array4_cupy():
+    # https://docs.cupy.dev/en/stable/user_guide/interoperability.html
+    import cupy as cp
+
+    # cupy -> AMReX Array4
+    x = np.ones(
+        (
+            2,
+            3,
+            4,
+        )
+    )  # TODO: merge into next line and create on device?
+    x_cupy = cp.asarray(x)  # type: cupy.ndarray
+    print(f"x_cupy={x_cupy}")
+    print(x_cupy.__cuda_array_interface__)
+
+    # cupy -> AMReX array4
+    x_arr = amrex.Array4_double(x_cupy)  # type: amrex.Array4_double
+    print(f"x_arr={x_arr}")
+    print(x_arr.__cuda_array_interface__)
+
+    assert (
+        x_arr.__cuda_array_interface__["data"][0]
+        == x_cupy.__cuda_array_interface__["data"][0]
+    )
+
+    # AMReX -> cupy
+    # arr_numba = cuda.as_cuda_array(arr4)
+    # ... or as MultiFab test
+    # TODO
+
+
+@pytest.mark.skipif(
+    amrex.Config.gpu_backend != "CUDA", reason="Requires AMReX_GPU_BACKEND=CUDA"
+)
+def test_array4_pytorch():
+    # https://docs.cupy.dev/en/stable/user_guide/interoperability.html#pytorch
+    # arr_torch = torch.as_tensor(arr, device='cuda')
+    # assert(arr_torch.__cuda_array_interface__['data'][0] == arr.__cuda_array_interface__['data'][0])
+    # TODO
+
+    pass