Merge pull request #9635 from Artemy-Mellanox/topic/multi_send-v5.0

pml/ucx: pml_ucx_multi_send_nb support

Author: janjust

config/ompi_check_ucx.m4

@@ -138,7 +138,8 @@ AC_DEFUN([OMPI_CHECK_UCX],[
                                          UCP_ATOMIC_FETCH_OP_FOR,
                                          UCP_ATOMIC_FETCH_OP_FXOR,
                                          UCP_PARAM_FIELD_ESTIMATED_NUM_PPN,
-                                         UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK],
+                                         UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK,
+                                         UCP_OP_ATTR_FLAG_MULTI_SEND],
                                         [], [],
                                         [#include <ucp/api/ucp.h>])
                          AC_CHECK_DECLS([UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS],

ompi/mca/pml/ucx/pml_ucx.c

@@ -566,7 +566,7 @@ int mca_pml_ucx_irecv(void *buf, size_t count, ompi_datatype_t *datatype,
 {
 #if HAVE_DECL_UCP_TAG_RECV_NBX
     pml_ucx_datatype_t *op_data = mca_pml_ucx_get_op_data(datatype);
-    ucp_request_param_t *param  = &op_data->op_param.recv;
+    ucp_request_param_t *param  = &op_data->op_param.irecv;
 #endif
 
     ucp_tag_t ucp_tag, ucp_tag_mask;
@@ -834,7 +834,7 @@ int mca_pml_ucx_isend(const void *buf, size_t count, ompi_datatype_t *datatype,
 #if HAVE_DECL_UCP_TAG_SEND_NBX
     req = (ompi_request_t*)mca_pml_ucx_common_send_nbx(ep, buf, count, datatype,
                                                        PML_UCX_MAKE_SEND_TAG(tag, comm), mode,
-                                                       &mca_pml_ucx_get_op_data(datatype)->op_param.send);
+                                                       &mca_pml_ucx_get_op_data(datatype)->op_param.isend);
 #else
     req = (ompi_request_t*)mca_pml_ucx_common_send(ep, buf, count, datatype,
                                                    mca_pml_ucx_get_datatype(datatype),

ompi/mca/pml/ucx/pml_ucx.h

@@ -59,6 +59,7 @@ struct mca_pml_ucx_module {
     int                       priority;
     bool                      cuda_initialized;
     bool                      request_leak_check;
+    uint32_t                  op_attr_nonblocking;
 };
 
 extern mca_pml_base_component_2_1_0_t mca_pml_ucx_component;

ompi/mca/pml/ucx/pml_ucx_component.c

@@ -49,6 +49,8 @@ mca_pml_base_component_2_1_0_t mca_pml_ucx_component = {
 
 static int mca_pml_ucx_component_register(void)
 {
+    int multi_send_op_attr_enable;
+
     ompi_pml_ucx.priority = 51;
     (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "priority",
                                            "Priority of the UCX component",
@@ -79,6 +81,20 @@ static int mca_pml_ucx_component_register(void)
     ompi_pml_ucx.request_leak_check = true;
 #endif
 
+    ompi_pml_ucx.op_attr_nonblocking = 0;
+#if HAVE_DECL_UCP_OP_ATTR_FLAG_MULTI_SEND
+    multi_send_op_attr_enable        = 0;
+    (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "multi_send_nb",
+                                           "Enable passing multi-send optimization flag for nonblocking operations",
+                                           MCA_BASE_VAR_TYPE_BOOL, NULL, 0, 0,
+                                           OPAL_INFO_LVL_3,
+                                           MCA_BASE_VAR_SCOPE_LOCAL,
+                                           &multi_send_op_attr_enable);
+    if (multi_send_op_attr_enable) {
+        ompi_pml_ucx.op_attr_nonblocking = UCP_OP_ATTR_FLAG_MULTI_SEND;
+    }
+#endif
+
     opal_common_ucx_mca_var_register(&mca_pml_ucx_component.pmlm_version);
     return 0;
 }

ompi/mca/pml/ucx/pml_ucx_datatype.c

@@ -24,7 +24,6 @@
 #ifdef HAVE_UCP_REQUEST_PARAM_T
 #define PML_UCX_DATATYPE_SET_VALUE(_datatype, _val) \
     (_datatype)->op_param.send._val; \
-    (_datatype)->op_param.bsend._val; \
     (_datatype)->op_param.recv._val;
 #endif
 
@@ -190,8 +189,6 @@ pml_ucx_datatype_t *mca_pml_ucx_init_nbx_datatype(ompi_datatype_t *datatype,
     pml_datatype->datatype                    = ucp_datatype;
     pml_datatype->op_param.send.op_attr_mask  = UCP_OP_ATTR_FIELD_CALLBACK;
     pml_datatype->op_param.send.cb.send       = mca_pml_ucx_send_nbx_completion;
-    pml_datatype->op_param.bsend.op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK;
-    pml_datatype->op_param.bsend.cb.send      = mca_pml_ucx_bsend_nbx_completion;
     pml_datatype->op_param.recv.op_attr_mask  = UCP_OP_ATTR_FIELD_CALLBACK |
                                                 UCP_OP_ATTR_FLAG_NO_IMM_CMPL;
     pml_datatype->op_param.recv.cb.recv       = mca_pml_ucx_recv_nbx_completion;
@@ -206,6 +203,11 @@ pml_ucx_datatype_t *mca_pml_ucx_init_nbx_datatype(ompi_datatype_t *datatype,
         PML_UCX_DATATYPE_SET_VALUE(pml_datatype, datatype = ucp_datatype);
     }
 
+    pml_datatype->op_param.isend = pml_datatype->op_param.send;
+    pml_datatype->op_param.irecv = pml_datatype->op_param.recv;
+    pml_datatype->op_param.isend.op_attr_mask |= ompi_pml_ucx.op_attr_nonblocking;
+    pml_datatype->op_param.irecv.op_attr_mask |= ompi_pml_ucx.op_attr_nonblocking;
+
     return pml_datatype;
 }
 #endif

ompi/mca/pml/ucx/pml_ucx_datatype.h

@@ -21,8 +21,9 @@ typedef struct {
     int                     size_shift;
     struct {
         ucp_request_param_t send;
-        ucp_request_param_t bsend;
+        ucp_request_param_t isend;
         ucp_request_param_t recv;
+        ucp_request_param_t irecv;
     } op_param;
 } pml_ucx_datatype_t;
 #endif