stream: add MPIX_I{send,recv}_enqueue and MPIX_Wait{all}_enqueue

pmodels · Apr 27, 2022 · d275aaa · d275aaa
1 parent 8c36f68
commit d275aaa
Show file tree

Hide file tree

Showing 4 changed files with 346 additions and 5 deletions.
diff --git a/maint/local_python/binding_c.py b/maint/local_python/binding_c.py
@@ -696,6 +696,9 @@ def process_func_parameters(func):
                 if RE.match(r'mpi_startall', func['name'], re.IGNORECASE):
                     impl_arg_list.append(ptrs_name)
                     impl_param_list.append("MPIR_Request **%s" % ptrs_name)
+                else:
+                    impl_arg_list.append(name)
+                    impl_param_list.append("MPI_Request %s[]" % name)
             else:
                 print("Unhandled handle array: " + name, file=sys.stderr)
         elif "code-handle_ptr-tail" in func and name in func['code-handle_ptr-tail']:

diff --git a/src/binding/c/stream_api.txt b/src/binding/c/stream_api.txt
@@ -28,3 +28,35 @@ MPIX_Recv_enqueue:
     tag: TAG, [message tag or MPI_ANY_TAG]
     comm: COMMUNICATOR
     status: STATUS, direction=out
+
+MPIX_Isend_enqueue:
+    buf: BUFFER, constant=True, [initial address of send buffer]
+    count: POLYXFER_NUM_ELEM_NNI, [number of elements in send buffer]
+    datatype: DATATYPE, [datatype of each send buffer element]
+    dest: RANK, [rank of destination]
+    tag: TAG, [message tag]
+    comm: COMMUNICATOR
+    request: REQUEST, direction=out
+
+MPIX_Irecv_enqueue:
+    buf: BUFFER, direction=out, [initial address of receive buffer]
+    count: POLYXFER_NUM_ELEM_NNI, [number of elements in receive buffer]
+    datatype: DATATYPE, [datatype of each receive buffer element]
+    source: RANK, [rank of source or MPI_ANY_SOURCE]
+    tag: TAG, [message tag or MPI_ANY_TAG]
+    comm: COMMUNICATOR
+    request: REQUEST, direction=out
+
+MPI_Wait_enqueue:
+    request: REQUEST, direction=inout, [request]
+    status: STATUS, direction=out
+
+MPI_Waitall_enqueue:
+    count: ARRAY_LENGTH_NNI, [lists length]
+    array_of_requests: REQUEST, direction=inout, length=count, [array of requests]
+    array_of_statuses: STATUS, direction=out, length=*, pointer=False, [array of status objects]
+{ -- error_check -- array_of_statuses
+    if (count > 0) {
+        MPIR_ERRTEST_ARGNULL(array_of_statuses, "array_of_statuses", mpi_errno);
+    }
+}
diff --git a/src/include/mpir_request.h b/src/include/mpir_request.h
@@ -65,6 +65,7 @@ typedef enum MPIR_Request_kind_t {
     MPIR_REQUEST_KIND__PART_SEND,       /* Partitioned send req returned to user */
     MPIR_REQUEST_KIND__PART_RECV,       /* Partitioned recv req returned to user */
     MPIR_REQUEST_KIND__PART,    /* Partitioned pt2pt internal reqs */
+    MPIR_REQUEST_KIND__ENQUEUE, /* enqueued (to gpu stream) request */
     MPIR_REQUEST_KIND__GREQUEST,
     MPIR_REQUEST_KIND__COLL,
     MPIR_REQUEST_KIND__MPROBE,  /* see NOTE-R1 */
@@ -220,6 +221,12 @@ struct MPIR_Request {
             MPL_atomic_int_t active_flag;       /* flag indicating whether in a start-complete active period.
                                                  * Value is 0 or 1. */
         } part;                 /* kind : MPIR_REQUEST_KIND__PART_SEND or MPIR_REQUEST_KIND__PART_RECV */
+        struct {
+            MPL_gpu_stream_t gpu_stream;
+            struct MPIR_Request *real_request;
+            bool is_send;
+            void *data;
+        } enqueue;
         struct {
             MPIR_Win *win;
         } rma;                  /* kind : MPIR_REQUEST_KIND__RMA */