REVIEW: address feedback

TL/UCP: update ep flush with cb

Author: wfaderhold21

src/components/tl/ucp/alltoall/alltoall_onesided.c

@@ -147,6 +147,7 @@ void ucc_tl_ucp_alltoall_onesided_put_progress(ucc_coll_task_t *ctask)
                               PTR_OFFSET(dest, grank * nelems), nelems,
                               peer, src_memh, dst_memh, team, task),
             task, out);
+        UCPCHECK_GOTO(ucc_tl_ucp_ep_flush(peer, team, task), task, out);
 
         if (!alltoall_onesided_handle_completion(task, posted, completed,
                                                  ntokens, npolls)) {
@@ -263,7 +264,7 @@ ucc_status_t ucc_tl_ucp_alltoall_onesided_init(ucc_base_coll_args_t *coll_args,
     nelems             = nelems / UCC_TL_TEAM_SIZE(tl_team);
     param.field_mask   = UCP_EP_PERF_PARAM_FIELD_MESSAGE_SIZE;
     attr.field_mask    = UCP_EP_PERF_ATTR_FIELD_ESTIMATED_TIME;
-    param.message_size = (1 << 20);
+    param.message_size = nelems * ucc_dt_size(TASK_ARGS(task).src.info.datatype);;
     ucc_tl_ucp_get_ep(
         tl_team, (UCC_TL_TEAM_RANK(tl_team) + 1) % UCC_TL_TEAM_SIZE(tl_team),
         &ep);

src/components/tl/ucp/tl_ucp_coll.h

@@ -201,7 +201,8 @@ static inline ucc_status_t ucc_tl_ucp_test_ring(ucc_tl_ucp_task_t *task)
 
 #define UCC_TL_UCP_TASK_ONESIDED_P2P_COMPLETE(_task)                           \
     (((_task)->onesided.put_posted == (_task)->onesided.put_completed) &&      \
-     ((_task)->onesided.get_posted == (_task)->onesided.get_completed))
+     ((_task)->onesided.get_posted == (_task)->onesided.get_completed) &&      \
+     ((_task)->flush_posted == (_task)->flush_completed))
 
 #define UCC_TL_UCP_TASK_ONESIDED_SYNC_COMPLETE(_task, _end)                    \
     (*((long *)(TASK_ARGS(_task).global_work_buffer)) == _end)

src/components/tl/ucp/tl_ucp_sendrecv.c

@@ -69,6 +69,19 @@ void ucc_tl_ucp_get_completion_cb(void *request, ucs_status_t status,
     ucp_request_free(request);
 }
 
+void ucc_tl_ucp_flush_completion_cb(void *request, ucs_status_t status,
+                                    void *user_data)
+{
+    ucc_tl_ucp_task_t *task = (ucc_tl_ucp_task_t *)user_data;
+    if (ucc_unlikely(UCS_OK != status)) {
+        tl_error(UCC_TASK_LIB(task), "failure in ep flush completion %s",
+                 ucs_status_string(status));
+        task->super.status = ucs_status_to_ucc_status(status);
+    }
+    task->flush_completed++;
+    ucp_request_free(request);
+}
+
 void ucc_tl_ucp_recv_completion_cb_mt(void *request, ucs_status_t status,
                                       const ucp_tag_recv_info_t *info, /* NOLINT */
                                       void *user_data)

src/components/tl/ucp/tl_ucp_sendrecv.h

@@ -62,6 +62,9 @@ void ucc_tl_ucp_put_completion_cb(void *request, ucs_status_t status,
 void ucc_tl_ucp_get_completion_cb(void *request, ucs_status_t status,
                                   void *user_data);
 
+void ucc_tl_ucp_flush_completion_cb(void *request, ucs_status_t status,
+                                    void *user_data);
+
 void ucc_tl_ucp_recv_completion_cb_st(void *request, ucs_status_t status,
                                       const ucp_tag_recv_info_t *info,
                                       void *user_data);
@@ -531,7 +534,8 @@ static inline ucc_status_t ucc_tl_ucp_flush(ucc_tl_ucp_team_t *team)
 }
 
 static inline ucc_status_t ucc_tl_ucp_ep_flush(ucc_rank_t dest_group_rank,
-                                               ucc_tl_ucp_team_t *team)
+                                               ucc_tl_ucp_team_t *team,
+                                               ucc_tl_ucp_task_t *task)
 {
     ucp_request_param_t req_param = {0};
     ucc_status_t        status;
@@ -543,12 +547,19 @@ static inline ucc_status_t ucc_tl_ucp_ep_flush(ucc_rank_t dest_group_rank,
         return status;
     }
 
+    req_param.op_attr_mask =
+        UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_USER_DATA;
+    req_param.cb.send   = ucc_tl_ucp_flush_completion_cb;
+    req_param.user_data = (void *)task;
+
     req = ucp_ep_flush_nbx(ep, &req_param);
+    task->flush_posted++;
     if (UCS_OK != req) {
         if (UCS_PTR_IS_ERR(req)) {
             return ucs_status_to_ucc_status(UCS_PTR_STATUS(req));
         }
-        ucp_request_free(req);
+    } else {
+        task->flush_completed++;
     }
     return UCC_OK;
 }

src/components/tl/ucp/tl_ucp_task.h

@@ -221,6 +221,8 @@ typedef struct ucc_tl_ucp_task {
         } alltoall_onesided;
         char                        plugin_data[UCC_TL_UCP_TASK_PLUGIN_MAX_DATA];
     };
+    uint32_t flush_posted;
+    uint32_t flush_completed;
 } ucc_tl_ucp_task_t;
 
 static inline void ucc_tl_ucp_task_reset(ucc_tl_ucp_task_t *task,
@@ -230,6 +232,8 @@ static inline void ucc_tl_ucp_task_reset(ucc_tl_ucp_task_t *task,
     task->tagged.send_completed = 0;
     task->tagged.recv_posted    = 0;
     task->tagged.recv_completed = 0;
+    task->flush_posted          = 0;
+    task->flush_completed       = 0;
     task->super.status          = status;
 }