diff --git a/src/ucp/core/ucp_am.c b/src/ucp/core/ucp_am.c index e2d2b095314..7d6228af55f 100644 --- a/src/ucp/core/ucp_am.c +++ b/src/ucp/core/ucp_am.c @@ -649,6 +649,8 @@ static ucs_status_t ucp_am_send_start_rndv(ucp_request_t *sreq) sreq->send.length); UCS_PROFILE_REQUEST_EVENT(sreq, "start_rndv", sreq->send.length); + sreq->send.msg_proto.sreq_id = ucp_send_request_get_id(sreq); + /* Note: no need to call ucp_ep_resolve_remote_id() here, because it * was done in ucp_am_send_nbx */ diff --git a/src/ucp/core/ucp_request.c b/src/ucp/core/ucp_request.c index 5afd8c0c722..668b34cd036 100644 --- a/src/ucp/core/ucp_request.c +++ b/src/ucp/core/ucp_request.c @@ -338,6 +338,14 @@ void ucp_request_init_multi_proto(ucp_request_t *req, UCS_PROFILE_REQUEST_EVENT(req, multi_func_str, req->send.length); } +static void ucp_send_request_init_id(ucp_request_t *req) +{ + if (req->flags & UCP_REQUEST_FLAG_SYNC) { + ucs_assert(req->flags & UCP_REQUEST_FLAG_SEND_TAG); + req->send.msg_proto.sreq_id = ucp_send_request_get_id(req); + } +} + ucs_status_t ucp_request_send_start(ucp_request_t *req, ssize_t max_short, size_t zcopy_thresh, size_t zcopy_max, @@ -366,6 +374,8 @@ ucp_request_send_start(ucp_request_t *req, ssize_t max_short, ucp_request_init_multi_proto(req, proto->bcopy_multi, "start_bcopy_multi"); } + + ucp_send_request_init_id(req); return UCS_OK; } else if (length < zcopy_max) { /* zcopy */ @@ -396,6 +406,8 @@ ucp_request_send_start(ucp_request_t *req, ssize_t max_short, req->send.uct.func = proto->zcopy_single; UCS_PROFILE_REQUEST_EVENT(req, "start_zcopy_single", req->send.length); } + + ucp_send_request_init_id(req); return UCS_OK; } diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index b08c1f1585b..0a98750831e 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -134,8 +134,11 @@ struct ucp_request { ucp_lane_index_t am_bw_index; /* AM BW lane index */ uint64_t message_id; /* used to identify matching parts of a large message */ + ucs_ptr_map_key_t sreq_id; /* send request ID on the + sender side */ ucs_ptr_map_key_t rreq_id; /* receive request ID on the recv side (used in AM rndv) */ + union { struct { ucp_tag_t tag; @@ -207,7 +210,7 @@ struct ucp_request { } rkey_ptr; struct { - ucs_ptr_map_key_t req_id; /* the send request ID on receiver side */ + ucs_ptr_map_key_t sreq_id; /* the send request ID on receiver side */ ucp_request_t *rreq; /* pointer to the receive request */ size_t length; /* the length of the data that should be fetched * from sender side */ @@ -295,6 +298,7 @@ struct ucp_request { uct_tag_context_t uct_ctx; /* Transport offload context */ ssize_t remaining; /* How much more data * to be received */ + ucs_ptr_map_key_t rreq_id; /* the receive request ID on receiver side */ union { struct { diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 50136bfe07f..9880fea4537 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -100,7 +100,7 @@ size_t ucp_rndv_rts_pack(ucp_request_t *sreq, ucp_rndv_rts_hdr_t *rndv_rts_hdr, void *rkey_buf; rndv_rts_hdr->sreq.ep_id = ucp_send_request_get_ep_remote_id(sreq); - rndv_rts_hdr->sreq.req_id = ucp_send_request_get_id(sreq); + rndv_rts_hdr->sreq.req_id = sreq->send.msg_proto.sreq_id; rndv_rts_hdr->size = sreq->send.length; rndv_rts_hdr->flags = flags; @@ -135,13 +135,11 @@ static size_t ucp_rndv_rtr_pack(void *dest, void *arg) ucp_request_t *rndv_req = arg; ucp_rndv_rtr_hdr_t *rndv_rtr_hdr = dest; ucp_request_t *rreq = rndv_req->send.rndv_rtr.rreq; - ucp_ep_h ep = rndv_req->send.ep; ssize_t packed_rkey_size; - rndv_rtr_hdr->sreq_id = rndv_req->send.rndv_rtr.req_id; + rndv_rtr_hdr->sreq_id = rndv_req->send.rndv_rtr.sreq_id; /* request of receiver side */ - rndv_rtr_hdr->rreq_id = ucp_worker_get_request_id(ep->worker, rreq, - ucp_ep_use_indirect_id(ep)); + rndv_rtr_hdr->rreq_id = rreq->recv.rreq_id; /* Pack remote keys (which can be empty list) */ if (UCP_DT_IS_CONTIG(rreq->recv.datatype)) { @@ -404,15 +402,21 @@ static void ucp_rndv_req_send_rtr(ucp_request_t *rndv_req, ucp_request_t *rreq, ucs_ptr_map_key_t sender_req_id, size_t recv_length, size_t offset) { + ucp_ep_h ep = rndv_req->send.ep; + ucp_trace_req(rndv_req, "send rtr remote sreq_id 0x%"PRIxPTR" rreq %p", sender_req_id, rreq); - rndv_req->send.lane = ucp_ep_get_am_lane(rndv_req->send.ep); - rndv_req->send.uct.func = ucp_proto_progress_rndv_rtr; - rndv_req->send.rndv_rtr.req_id = sender_req_id; - rndv_req->send.rndv_rtr.rreq = rreq; - rndv_req->send.rndv_rtr.length = recv_length; - rndv_req->send.rndv_rtr.offset = offset; + rndv_req->send.lane = ucp_ep_get_am_lane(rndv_req->send.ep); + rndv_req->send.uct.func = ucp_proto_progress_rndv_rtr; + rndv_req->send.rndv_rtr.sreq_id = sender_req_id; + rndv_req->send.rndv_rtr.rreq = rreq; + rndv_req->send.rndv_rtr.length = recv_length; + rndv_req->send.rndv_rtr.offset = offset; + + rndv_req->send.rndv_rtr.rreq->recv.rreq_id = + ucp_worker_get_request_id(ep->worker, rreq, + ucp_ep_use_indirect_id(ep)); ucp_request_send(rndv_req, 0); } @@ -1668,6 +1672,9 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler, /* Do not deregister memory here, because am zcopy rndv may * need it registered (if am and tag is the same lane). */ ucp_tag_offload_cancel_rndv(sreq); + ucs_assert(!ucp_ep_use_indirect_id(ep)); + } else { + ucp_worker_del_request_id(sreq->send.ep->worker, sreq->send.msg_proto.sreq_id); } if (UCP_DT_IS_CONTIG(sreq->send.datatype) && rndv_rtr_hdr->address) { diff --git a/src/ucp/tag/eager_snd.c b/src/ucp/tag/eager_snd.c index 110af2030e7..7423ec6dc7a 100644 --- a/src/ucp/tag/eager_snd.c +++ b/src/ucp/tag/eager_snd.c @@ -52,7 +52,7 @@ static size_t ucp_tag_pack_eager_sync_only_dt(void *dest, void *arg) hdr->super.super.tag = req->send.msg_proto.tag.tag; hdr->req.ep_id = ucp_send_request_get_ep_remote_id(req); - hdr->req.req_id = ucp_send_request_get_id(req); + hdr->req.req_id = req->send.msg_proto.sreq_id; return ucp_tag_pack_eager_common(req, hdr + 1, req->send.length, sizeof(*hdr), 1); @@ -92,7 +92,7 @@ static size_t ucp_tag_pack_eager_sync_first_dt(void *dest, void *arg) hdr->super.total_len = req->send.length; hdr->req.ep_id = ucp_send_request_get_ep_remote_id(req); hdr->super.msg_id = req->send.msg_proto.message_id; - hdr->req.req_id = ucp_send_request_get_id(req); + hdr->req.req_id = req->send.msg_proto.sreq_id; return ucp_tag_pack_eager_common(req, hdr + 1, length, sizeof(*hdr), 1); } @@ -256,7 +256,7 @@ static ucs_status_t ucp_tag_eager_sync_zcopy_single(uct_pending_req_t *self) hdr.super.super.tag = req->send.msg_proto.tag.tag; hdr.req.ep_id = ucp_send_request_get_ep_remote_id(req); - hdr.req.req_id = ucp_send_request_get_id(req); + hdr.req.req_id = req->send.msg_proto.sreq_id; return ucp_do_am_zcopy_single(self, UCP_AM_ID_EAGER_SYNC_ONLY, &hdr, sizeof(hdr), NULL, 0ul, @@ -272,7 +272,7 @@ static ucs_status_t ucp_tag_eager_sync_zcopy_multi(uct_pending_req_t *self) first_hdr.super.super.super.tag = req->send.msg_proto.tag.tag; first_hdr.super.total_len = req->send.length; first_hdr.req.ep_id = ucp_send_request_get_ep_remote_id(req); - first_hdr.req.req_id = ucp_send_request_get_id(req); + first_hdr.req.req_id = req->send.msg_proto.sreq_id; first_hdr.super.msg_id = req->send.msg_proto.message_id; middle_hdr.msg_id = req->send.msg_proto.message_id; middle_hdr.offset = req->send.state.dt.offset; diff --git a/src/ucp/tag/offload.c b/src/ucp/tag/offload.c index 5a89d82f8ce..3297e44be09 100644 --- a/src/ucp/tag/offload.c +++ b/src/ucp/tag/offload.c @@ -560,6 +560,8 @@ static void ucp_tag_offload_rndv_zcopy_completion(uct_completion_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.state.uct_comp); + + ucs_assert(!ucp_ep_use_indirect_id(req->send.ep)); ucp_proto_am_zcopy_req_complete(req, self->status); } @@ -578,10 +580,11 @@ ucs_status_t ucp_tag_offload_rndv_zcopy(uct_pending_req_t *self) ucp_tag_offload_unexp_rndv_hdr_t rndv_hdr = { .ep_id = ucp_send_request_get_ep_remote_id(req), - .req_id = ucp_send_request_get_id(req), + .req_id = req->send.msg_proto.sreq_id, .md_index = md_index }; + ucs_assert(!ucp_ep_use_indirect_id(req->send.ep)); dt_state = req->send.state.dt; UCS_STATIC_ASSERT(sizeof(ucp_rsc_index_t) <= sizeof(rndv_hdr.md_index)); diff --git a/src/ucp/tag/tag_rndv.c b/src/ucp/tag/tag_rndv.c index 6a9779f183b..b2b8b146c48 100644 --- a/src/ucp/tag/tag_rndv.c +++ b/src/ucp/tag/tag_rndv.c @@ -105,6 +105,8 @@ ucs_status_t ucp_tag_send_start_rndv(ucp_request_t *sreq) return status; } + sreq->send.msg_proto.sreq_id = ucp_send_request_get_id(sreq); + if (ucp_ep_is_tag_offload_enabled(ucp_ep_config(ep))) { status = ucp_tag_offload_start_rndv(sreq); } else { diff --git a/src/ucs/datastruct/ptr_map.inl b/src/ucs/datastruct/ptr_map.inl index 0419d62a9b8..11920639ca2 100644 --- a/src/ucs/datastruct/ptr_map.inl +++ b/src/ucs/datastruct/ptr_map.inl @@ -79,6 +79,7 @@ ucs_ptr_map_put(ucs_ptr_map_t *map, void *ptr, int indirect, if (ucs_likely(!indirect)) { *key = (uintptr_t)ptr; ucs_assert(!(*key & UCS_PTR_MAP_KEY_MIN_ALIGN)); + ucs_assert(*key != 0); return UCS_OK; } @@ -131,6 +132,7 @@ ucs_ptr_map_extract(ucs_ptr_map_t *map, ucs_ptr_map_key_t key) void *value; if (ucs_likely(!(key & UCS_PTR_MAP_KEY_INDIRECT_FLAG))) { + ucs_assert(key != 0); return (void*)key; } diff --git a/test/gtest/ucp/test_ucp_sockaddr.cc b/test/gtest/ucp/test_ucp_sockaddr.cc index 018403ed2ad..782d8104839 100644 --- a/test/gtest/ucp/test_ucp_sockaddr.cc +++ b/test/gtest/ucp/test_ucp_sockaddr.cc @@ -978,34 +978,38 @@ class test_ucp_sockaddr_protocols : public test_ucp_sockaddr { void test_tag_send_recv(size_t size, bool is_exp, bool is_sync = false) { - std::string send_buf(size, 'x'); - std::string recv_buf(size, 'y'); + /* send multiple messages to test the protocol both before and after + * connection establishment */ + for (int i = 0; i < 100; i++) { + std::string send_buf(size, 'x'); + std::string recv_buf(size, 'y'); - void *rreq = NULL, *sreq = NULL; + void *rreq = NULL, *sreq = NULL; - if (is_exp) { - rreq = ucp_tag_recv_nb(receiver().worker(), &recv_buf[0], size, - ucp_dt_make_contig(1), 0, 0, rtag_complete_cb); - } + if (is_exp) { + rreq = ucp_tag_recv_nb(receiver().worker(), &recv_buf[0], size, + ucp_dt_make_contig(1), 0, 0, rtag_complete_cb); + } - if (is_sync) { - sreq = ucp_tag_send_sync_nb(sender().ep(), &send_buf[0], size, - ucp_dt_make_contig(1), 0, scomplete_cb); - } else { - sreq = ucp_tag_send_nb(sender().ep(), &send_buf[0], size, - ucp_dt_make_contig(1), 0, scomplete_cb); - } + if (is_sync) { + sreq = ucp_tag_send_sync_nb(sender().ep(), &send_buf[0], size, + ucp_dt_make_contig(1), 0, scomplete_cb); + } else { + sreq = ucp_tag_send_nb(sender().ep(), &send_buf[0], size, + ucp_dt_make_contig(1), 0, scomplete_cb); + } - if (!is_exp) { - short_progress_loop(); - rreq = ucp_tag_recv_nb(receiver().worker(), &recv_buf[0], size, - ucp_dt_make_contig(1), 0, 0, rtag_complete_cb); - } + if (!is_exp) { + short_progress_loop(); + rreq = ucp_tag_recv_nb(receiver().worker(), &recv_buf[0], size, + ucp_dt_make_contig(1), 0, 0, rtag_complete_cb); + } - request_wait(sreq); - request_wait(rreq); + request_wait(sreq); + request_wait(rreq); - compare_buffers(send_buf, recv_buf); + compare_buffers(send_buf, recv_buf); + } } void wait_for_server_ep()