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_iprobe(int src, int tag, struct ompi_communicator_t* comm,
int *matched, ompi_status_public_t* mpi_status)
{
ucp_tag_t ucp_tag, ucp_tag_mask;
ucp_tag_recv_info_t info;
ucp_tag_message_h ucp_msg;
PML_UCX_TRACE_PROBE("iprobe", src, tag, comm);
PML_UCX_MAKE_RECV_TAG(ucp_tag, ucp_tag_mask, tag, src, comm);
ucp_msg = ucp_tag_probe_nb(ompi_pml_ucx.ucp_worker, ucp_tag, ucp_tag_mask,
0, &info);
if (ucp_msg != NULL) {
*matched = 1;
mca_pml_ucx_set_recv_status_safe(mpi_status, UCS_OK, &info);
} else {
*matched = 0;
}
return OMPI_SUCCESS;
}
int mca_pml_ucx_improbe(int src, int tag, struct ompi_communicator_t* comm,
int *matched, struct ompi_message_t **message,
ompi_status_public_t* mpi_status)
{
ucp_tag_t ucp_tag, ucp_tag_mask;
ucp_tag_recv_info_t info;
ucp_tag_message_h ucp_msg;
PML_UCX_TRACE_PROBE("improbe", src, tag, comm);
PML_UCX_MAKE_RECV_TAG(ucp_tag, ucp_tag_mask, tag, src, comm);
ucp_msg = ucp_tag_probe_nb(ompi_pml_ucx.ucp_worker, ucp_tag, ucp_tag_mask,
1, &info);
if (ucp_msg != NULL) {
PML_UCX_MESSAGE_NEW(comm, ucp_msg, &info, message);
PML_UCX_VERBOSE(8, "got message %p (%p)", (void*)*message, (void*)ucp_msg);
*matched = 1;
mca_pml_ucx_set_recv_status_safe(mpi_status, UCS_OK, &info);
} else if (UCS_PTR_STATUS(ucp_msg) == UCS_ERR_NO_MESSAGE) {
*matched = 0;
}
return OMPI_SUCCESS;
}
int mca_pml_ucx_mprobe(int src, int tag, struct ompi_communicator_t* comm,
struct ompi_message_t **message,
ompi_status_public_t* mpi_status)
{
ucp_tag_t ucp_tag, ucp_tag_mask;
ucp_tag_recv_info_t info;
ucp_tag_message_h ucp_msg;
PML_UCX_TRACE_PROBE("mprobe", src, tag, comm);
PML_UCX_MAKE_RECV_TAG(ucp_tag, ucp_tag_mask, tag, src, comm);
for (;;) {
ucp_msg = ucp_tag_probe_nb(ompi_pml_ucx.ucp_worker, ucp_tag, ucp_tag_mask,
1, &info);
if (ucp_msg != NULL) {
PML_UCX_MESSAGE_NEW(comm, ucp_msg, &info, message);
PML_UCX_VERBOSE(8, "got message %p (%p)", (void*)*message, (void*)ucp_msg);
mca_pml_ucx_set_recv_status_safe(mpi_status, UCS_OK, &info);
return OMPI_SUCCESS;
}
opal_progress();
}
}
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();
}
}
