struct ucx_context * launch_recv() { ucp_tag_recv_info_t info_tag; ucp_tag_message_h msg_tag; ucs_status_t status; struct ucx_context *request; static char *msg = NULL; static int cur_len = 0; msg_tag = ucp_tag_probe_nb(ucp_worker, tag, tag_mask, 1, &info_tag); if(msg_tag == NULL) { return NULL; } if( same_buf ) { if( cur_len < info_tag.length ){ if( NULL == msg ) { msg = malloc(info_tag.length); } else { free(msg); msg = malloc(info_tag.length); } cur_len = info_tag.length; } } else { msg = malloc(info_tag.length); } if (!msg) { fprintf(stderr, "unable to allocate memory\n"); abort(); } request = ucp_tag_msg_recv_nb(ucp_worker, msg, info_tag.length, ucp_dt_make_contig(1), msg_tag, recv_handle); if (UCS_PTR_IS_ERR(request)) { fprintf(stderr, "unable to receive UCX data message (%u)\n", UCS_PTR_STATUS(request)); free(msg); abort(); } request->buf = msg; return request; }
int mca_pml_ucx_mrecv(void *buf, size_t count, ompi_datatype_t *datatype, struct ompi_message_t **message, ompi_status_public_t* status) { ompi_request_t *req; PML_UCX_TRACE_MRECV("mrecv", buf, count, datatype, message); req = (ompi_request_t*)ucp_tag_msg_recv_nb(ompi_pml_ucx.ucp_worker, buf, count, mca_pml_ucx_get_datatype(datatype), (*message)->req_ptr, mca_pml_ucx_recv_completion); if (UCS_PTR_IS_ERR(req)) { PML_UCX_ERROR("ucx msg recv failed: %s", ucs_status_string(UCS_PTR_STATUS(req))); return OMPI_ERROR; } PML_UCX_MESSAGE_RELEASE(message); ompi_request_wait(&req, status); return OMPI_SUCCESS; }
static int run_ucx_server(ucp_worker_h ucp_worker) { ucp_tag_recv_info_t info_tag; ucp_tag_message_h msg_tag; ucs_status_t status; ucp_ep_h client_ep; struct msg *msg = 0; struct ucx_context *request = 0; size_t msg_len = 0; int ret = -1; /* Receive client UCX address */ do { /* Following blocked methods used to polling internal file descriptor * to make CPU idle and don't spin loop */ if (ucp_test_mode == TEST_MODE_WAIT) { status = ucp_worker_wait(ucp_worker); if (status != UCS_OK) { goto err; } } else if (ucp_test_mode == TEST_MODE_EVENTFD) { status = test_poll_wait(ucp_worker); if (status != UCS_OK) { goto err; } } /* Progressing before probe to update the state */ ucp_worker_progress(ucp_worker); /* Probing incoming events in non-block mode */ msg_tag = ucp_tag_probe_nb(ucp_worker, tag, tag_mask, 1, &info_tag); } while (msg_tag == NULL); msg = malloc(info_tag.length); if (!msg) { fprintf(stderr, "unable to allocate memory\n"); goto err; } request = ucp_tag_msg_recv_nb(ucp_worker, msg, info_tag.length, ucp_dt_make_contig(1), msg_tag, recv_handle); if (UCS_PTR_IS_ERR(request)) { fprintf(stderr, "unable to receive UCX address message (%s)\n", ucs_status_string(UCS_PTR_STATUS(request))); free(msg); goto err; } else { wait(ucp_worker, request); ucp_request_release(request); printf("UCX address message was received\n"); } peer_addr = malloc(msg->data_len); if (!peer_addr) { fprintf(stderr, "unable to allocate memory for peer address\n"); free(msg); goto err; } peer_addr_len = msg->data_len; memcpy(peer_addr, msg->data, peer_addr_len); free(msg); /* Send test string to client */ status = ucp_ep_create(ucp_worker, peer_addr, &client_ep); if (status != UCS_OK) { goto err; } msg_len = sizeof(*msg) + strlen(test_str) + 1; msg = calloc(1, msg_len); if (!msg) { printf("unable to allocate memory\n"); goto err_ep; } msg->data_len = msg_len - sizeof(*msg); snprintf((char *)msg->data, msg->data_len, "%s", test_str); request = ucp_tag_send_nb(client_ep, msg, msg_len, ucp_dt_make_contig(1), tag, send_handle); if (UCS_PTR_IS_ERR(request)) { fprintf(stderr, "unable to send UCX data message\n"); free(msg); goto err_ep; } else if (UCS_PTR_STATUS(request) != UCS_OK) { printf("UCX data message was scheduled for send\n"); wait(ucp_worker, request); ucp_request_release(request); } ret = 0; free(msg); err_ep: ucp_ep_destroy(client_ep); err: return ret; }
static int connect_server() { ucp_tag_recv_info_t info_tag; ucp_tag_message_h msg_tag; ucs_status_t status; ucp_ep_params_t ep_params; struct msg *msg = 0; struct ucx_context *request = 0; /* Receive client UCX address */ do { /* Following blocked methods used to polling internal file descriptor * to make CPU idle and don't spin loop */ /* Progressing before probe to update the state */ ucp_worker_progress(ucp_worker); /* Probing incoming events in non-block mode */ msg_tag = ucp_tag_probe_nb(ucp_worker, tag, tag_mask, 1, &info_tag); } while (msg_tag == NULL); msg = malloc(info_tag.length); if (!msg) { fprintf(stderr, "unable to allocate memory\n"); abort(); } request = ucp_tag_msg_recv_nb(ucp_worker, msg, info_tag.length, ucp_dt_make_contig(1), msg_tag, recv_handle); if (UCS_PTR_IS_ERR(request)) { fprintf(stderr, "unable to receive UCX address message (%s)\n", ucs_status_string(UCS_PTR_STATUS(request))); free(msg); abort(); } else { wait(ucp_worker, request); request->completed = 0; ucp_request_release(request); printf("UCX address message was received\n"); } peer_addr = malloc(msg->data_len); if (!peer_addr) { fprintf(stderr, "unable to allocate memory for peer address\n"); free(msg); abort(); } peer_addr_len = msg->data_len; memcpy(peer_addr, msg->data, peer_addr_len); free(msg); /* Send test string to client */ ep_params.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS; ep_params.address = peer_addr; status = ucp_ep_create(ucp_worker, &ep_params, &rem_ep); if (status != UCS_OK) { abort(); } }