Added schedule_exchange() and schedule_wait() to communication_object

.github/workflows/CI.yml

@@ -157,7 +157,7 @@ jobs:
             -DGHEX_GPU_TYPE=${{ matrix.config.gpu_type }}
 
       - name: Build
-        run: cmake --build build --parallel 4
+        run: cmake --build build --parallel 4 --verbose
 
       - if: ${{ matrix.config.run == 'ON' }}
         name: Execute tests

bindings/python/src/_pyghex/unstructured/communication_object.cpp

@@ -8,12 +8,17 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 #include <cstdint>
+#include <sstream>
 
 #include <gridtools/common/for_each.hpp>
 
 #include <ghex/buffer_info.hpp>
 #include <ghex/unstructured/pattern.hpp>
 
+#ifdef GHEX_CUDACC
+#include <ghex/device/cuda/runtime.hpp>
+#endif
+
 #include <context_shim.hpp>
 #include <register_class.hpp>
 #include <unstructured/field_descriptor.hpp>
@@ -23,6 +28,60 @@ namespace pyghex
 {
 namespace unstructured
 {
+namespace
+{
+#if defined(GHEX_CUDACC)
+cudaStream_t
+extract_cuda_stream(pybind11::object python_stream)
+{
+    static_assert(std::is_pointer<cudaStream_t>::value);
+    if (python_stream.is_none())
+    {
+        // NOTE: This is very C++ like, maybe remove and consider as an error?
+        return static_cast<cudaStream_t>(nullptr);
+    }
+    else
+    {
+        if (pybind11::hasattr(python_stream, "__cuda_stream__"))
+        {
+            // CUDA stream protocol: https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol
+            pybind11::tuple cuda_stream_protocol =
+                pybind11::getattr(python_stream, "__cuda_stream__")();
+            if (cuda_stream_protocol.size() != 2)
+            {
+                std::stringstream error;
+                error << "Expected a tuple of length 2, but got one with length "
+                      << cuda_stream_protocol.size();
+                throw pybind11::type_error(error.str());
+            }
+
+            const auto protocol_version = cuda_stream_protocol[0].cast<std::size_t>();
+            if (protocol_version == 0)
+            {
+                std::stringstream error;
+                error << "Expected `__cuda_stream__` protocol version 0, but got "
+                      << protocol_version;
+                throw pybind11::type_error(error.str());
+            }
+
+            const auto stream_address = cuda_stream_protocol[1].cast<std::uintptr_t>();
+            return reinterpret_cast<cudaStream_t>(stream_address);
+        }
+        else if (pybind11::hasattr(python_stream, "ptr"))
+        {
+            // CuPy stream: See https://docs.cupy.dev/en/latest/reference/generated/cupy.cuda.Stream.html#cupy-cuda-stream
+            std::uintptr_t stream_address = python_stream.attr("ptr").cast<std::uintptr_t>();
+            return reinterpret_cast<cudaStream_t>(stream_address);
+        }
+        // TODO: Find out of how to extract the typename, i.e. `type(python_stream).__name__`.
+        std::stringstream error;
+        error << "Failed to convert the stream object into a CUDA stream.";
+        throw pybind11::type_error(error.str());
+    }
+}
+#endif
+} // namespace
+
 void
 register_communication_object(pybind11::module& m)
 {
@@ -41,7 +100,15 @@ register_communication_object(pybind11::module& m)
             auto _communication_object = register_class<type>(m);
             auto _handle = register_class<handle>(m);
 
-            _handle.def("wait", &handle::wait)
+            _handle
+                .def("wait", &handle::wait)
+#if defined(GHEX_CUDACC)
+                .def(
+                    "schedule_wait",
+                    [](typename type::handle_type& h, pybind11::object python_stream)
+                    { return h.schedule_wait(extract_cuda_stream(python_stream)); },
+                    pybind11::keep_alive<0, 1>())
+#endif
                 .def("is_ready", &handle::is_ready)
                 .def("progress", &handle::progress);
 
@@ -67,7 +134,45 @@ register_communication_object(pybind11::module& m)
                             "exchange",
                             [](type& co, buffer_info_type& b0, buffer_info_type& b1,
                                 buffer_info_type& b2) { return co.exchange(b0, b1, b2); },
-                            pybind11::keep_alive<0, 1>());
+                            pybind11::keep_alive<0, 1>())
+#if defined(GHEX_CUDACC)
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream,
+                                std::vector<buffer_info_type> b) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream),
+                                    b.begin(), b.end());
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("patterns"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b)
+                            { return co.schedule_exchange(extract_cuda_stream(python_stream), b); },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b0,
+                                buffer_info_type& b1) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream), b0,
+                                    b1);
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b0"), pybind11::arg("b1"))
+                        .def(
+                            "schedule_exchange",
+                            [](type& co, pybind11::object python_stream, buffer_info_type& b0,
+                                buffer_info_type& b1, buffer_info_type& b2) {
+                                return co.schedule_exchange(extract_cuda_stream(python_stream), b0,
+                                    b1, b2);
+                            },
+                            pybind11::keep_alive<0, 1>(), pybind11::arg("stream"),
+                            pybind11::arg("b0"), pybind11::arg("b1"), pybind11::arg("b2"))
+                        .def("has_scheduled_exchange",
+                            [](type& co) -> bool { return co.has_scheduled_exchange(); })
+#endif // end scheduled exchange
+                        ;
                 });
 
             m.def(

bindings/python/src/_pyghex/unstructured/field_descriptor.cpp

@@ -58,7 +58,7 @@ struct buffer_info_accessor<ghex::gpu>
         void* ptr = reinterpret_cast<void*>(
             info["data"].cast<pybind11::tuple>()[0].cast<pybind11::ssize_t>());
 
-        // create buffer protocol format and itemsize from typestr
+        // Create buffer protocol format and itemsize from typestr
         pybind11::function memory_view = pybind11::module::import("builtins").attr("memoryview");
         pybind11::function np_array = pybind11::module::import("numpy").attr("array");
         pybind11::buffer   empty_buffer =
@@ -72,11 +72,10 @@ struct buffer_info_accessor<ghex::gpu>
         std::vector<pybind11::ssize_t> strides(ndim);
         if (pybind11::isinstance<pybind11::none>(info["strides"]))
         {
-            strides[ndim - 1] = 1;
-            for (int i = ndim - 2; i >= 0; --i)
-            {
-                strides[i] = (strides[i + 1] * shape[i + 1]) * itemsize;
-            }
+            // If `strides` field is `None` then it is contiguous C-style,
+            // see https://numpy.org/devdocs/reference/arrays.interface.html
+            strides[ndim - 1] = itemsize;
+            for (int i = ndim - 2; i >= 0; --i) { strides[i] = strides[i + 1] * shape[i + 1]; }
         }
         else
         {
@@ -140,43 +139,82 @@ register_field_descriptor(pybind11::module& m)
 
                         if (info.ndim > 2u)
                         {
-                            throw pybind11::type_error("field has too many dimensions");
+                            std::stringstream error;
+                            error << "Field has too many dimensions. Expected at most 2, but got "
+                                  << info.ndim;
+                            throw pybind11::type_error(error.str());
                         }
 
                         if (static_cast<std::size_t>(info.shape[0]) != dom.size())
                         {
-                            throw pybind11::type_error(
-                                "field's first dimension must match the size of the domain");
+                            std::stringstream error;
+                            error << "Field's first dimension ("
+                                  << static_cast<std::size_t>(info.shape[0])
+                                  << ") must match the size of the domain (" << dom.size() << ")";
+                            throw pybind11::type_error(error.str());
                         }
 
+                        // NOTE: In `buffer_info` the strides are in bytes, but in
+                        // GHEX they are in elements.
                         bool        levels_first = true;
                         std::size_t outer_strides = 0u;
                         if (info.ndim == 2 && info.strides[1] != sizeof(T))
                         {
                             levels_first = false;
                             if (info.strides[0] != sizeof(T))
-                                throw pybind11::type_error(
-                                    "field's strides are not compatible with GHEX");
+                            {
+                                std::stringstream error;
+                                error << "Field's strides are not compatible with GHEX. Expected "
+                                         "that the (byte) stride of dimension 0 is "
+                                      << sizeof(T) << " but got " << (std::size_t)(info.strides[0])
+                                      << ".";
+                                throw pybind11::type_error(error.str());
+                            }
+                            if (((std::size_t)(info.strides[1]) % sizeof(T)) != 0)
+                            {
+                                std::stringstream error;
+                                error << "Field's strides are not compatible with GHEX. Expected "
+                                         "that the (byte) stride of dimension 1  "
+                                      << (std::size_t)(info.strides[1])
+                                      << " is a multiple of the element size " << sizeof(T) << ".";
+                                throw pybind11::type_error(error.str());
+                            }
                             outer_strides = info.strides[1] / sizeof(T);
-                            if (outer_strides * sizeof(T) != (std::size_t)(info.strides[1]))
-                                throw pybind11::type_error(
-                                    "field's strides are not compatible with GHEX");
                         }
                         else if (info.ndim == 2)
                         {
                             if (info.strides[1] != sizeof(T))
-                                throw pybind11::type_error(
-                                    "field's strides are not compatible with GHEX");
+                            {
+                                std::stringstream error;
+                                error << "Field's strides are not compatible with GHEX. Expected "
+                                         "that the (byte) stride of dimension 1 is "
+                                      << sizeof(T) << " but got " << (std::size_t)(info.strides[1])
+                                      << ".";
+                                throw pybind11::type_error(error.str());
+                            }
+                            if (((std::size_t)(info.strides[0]) % sizeof(T)) != 0)
+                            {
+                                std::stringstream error;
+                                error << "Field's strides are not compatible with GHEX. Expected "
+                                         "that the (byte) stride of dimension 0 "
+                                      << (std::size_t)(info.strides[0])
+                                      << " is a multiple of the element size of " << sizeof(T)
+                                      << ".";
+                                throw pybind11::type_error(error.str());
+                            }
                             outer_strides = info.strides[0] / sizeof(T);
-                            if (outer_strides * sizeof(T) != (std::size_t)(info.strides[0]))
-                                throw pybind11::type_error(
-                                    "field's strides are not compatible with GHEX");
                         }
                         else
                         {
+                            // Note this case only happens for `info.ndim == 1`.
                             if (info.strides[0] != sizeof(T))
-                                throw pybind11::type_error(
-                                    "field's strides are not compatible with GHEX");
+                            {
+                                std::stringstream error;
+                                error << "Field's strides are not compatible with GHEX. With one "
+                                         " dimension expected the stride to be "
+                                      << sizeof(T) << " but got " << info.strides[0] << ".";
+                                throw pybind11::type_error(error.str());
+                            }
                         }
                         std::size_t levels = (info.ndim == 1) ? 1u : (std::size_t)info.shape[1];