diff --git a/src/ucp/dt/datatype_iter.c b/src/ucp/dt/datatype_iter.c index c68893a8870..b1743bfcb54 100644 --- a/src/ucp/dt/datatype_iter.c +++ b/src/ucp/dt/datatype_iter.c @@ -357,13 +357,16 @@ ucs_status_t ucp_datatype_iter_sgl_mem_reg(ucp_context_h context, ucp_mem_h *memhs; size_t i; - if ((md_map == 0) || (dt_iter->type.sgl.memhs != NULL)) { + if (md_map == 0) { return UCS_OK; } - memhs = ucs_calloc(count, sizeof(*memhs), "dt_sgl_memh"); + memhs = dt_iter->type.sgl.memhs; if (memhs == NULL) { - return UCS_ERR_NO_MEMORY; + memhs = ucs_calloc(count, sizeof(*memhs), "dt_sgl_memh"); + if (memhs == NULL) { + return UCS_ERR_NO_MEMORY; + } } for (i = 0; i < count; ++i) { @@ -372,10 +375,14 @@ ucs_status_t ucp_datatype_iter_sgl_mem_reg(ucp_context_h context, dt_iter->type.sgl.lengths[i], dt_iter->mem_info.type, md_map, uct_flags, &memhs[i]); if (status != UCS_OK) { - while (i-- > 0) { - ucp_datatype_iter_mem_dereg_single(&memhs[i]); + if (dt_iter->type.sgl.memhs == NULL) { + /* Roll back only handles we own */ + while (i-- > 0) { + ucp_datatype_iter_mem_dereg_single(&memhs[i]); + } + ucs_free(memhs); } - ucs_free(memhs); + return status; } } diff --git a/src/ucp/rma/put_offload.c b/src/ucp/rma/put_offload.c index c06541120b0..abb3652ba8c 100644 --- a/src/ucp/rma/put_offload.c +++ b/src/ucp/rma/put_offload.c @@ -400,9 +400,8 @@ ucp_proto_put_sgl_offload_send_frag(ucp_request_t *req, size_t max_frag_length, ucp_datatype_iter_t *next_iter) { - ucp_ep_t *ep = req->send.ep; ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter; - ucp_md_index_t md_index = ucp_ep_md_index(ep, lpriv->super.lane); + ucp_rsc_index_t md_index = lpriv->super.md_index; ucp_mem_h *sgl_memhs = dt_iter->type.sgl.memhs; void *buffer = NULL; size_t length = 0; @@ -419,8 +418,10 @@ ucp_proto_put_sgl_offload_send_frag(ucp_request_t *req, return UCS_OK; } - uct_memh = (sgl_memhs != NULL) ? sgl_memhs[elem_index]->uct[md_index] : - UCT_MEM_HANDLE_NULL; + uct_memh = (sgl_memhs != NULL) ? + ucp_datatype_iter_uct_memh(sgl_memhs[elem_index], + md_index) : + UCT_MEM_HANDLE_NULL; uct_rkey = ucp_rkey_get_tl_rkey(req->send.rma.sgl.rkeys[elem_index], lpriv->super.rkey_index); @@ -435,9 +436,8 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req, ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift) { - ucp_ep_t *ep = req->send.ep; ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter; - ucp_md_index_t md_index = ucp_ep_md_index(ep, lpriv->super.lane); + ucp_rsc_index_t md_index = lpriv->super.md_index; ucp_rsc_index_t rkey_index = lpriv->super.rkey_index; size_t max_frag_length = lpriv->max_frag; ucp_mem_h *sgl_memhs = dt_iter->type.sgl.memhs; @@ -487,7 +487,8 @@ ucp_proto_put_sgl_offload_send_func(ucp_request_t *req, } uct_memhs[elem_count] = (sgl_memhs != NULL) ? - sgl_memhs[idx]->uct[md_index] : + ucp_datatype_iter_uct_memh(sgl_memhs[idx], + md_index) : UCT_MEM_HANDLE_NULL; uct_rkeys[elem_count] = ucp_rkey_get_tl_rkey(sgl_rkeys[idx], rkey_index); @@ -549,7 +550,7 @@ ucp_proto_put_sgl_offload_sw_send_func(ucp_request_t *req, ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter; ucp_lane_index_t lane = lpriv->super.lane; uct_ep_h uct_ep = ucp_ep_get_lane(ep, lane); - ucp_md_index_t md_index = ucp_ep_md_index(ep, lane); + ucp_rsc_index_t md_index = lpriv->super.md_index; ucp_rsc_index_t rkey_index = lpriv->super.rkey_index; size_t max_frag_length = lpriv->max_frag; ucp_mem_h *sgl_memhs = dt_iter->type.sgl.memhs; @@ -581,8 +582,10 @@ ucp_proto_put_sgl_offload_sw_send_func(ucp_request_t *req, tl_rkey = ucp_rkey_get_tl_rkey(sgl_rkeys[elem_index], rkey_index); iov.buffer = buffer; iov.length = length; - iov.memh = (sgl_memhs != NULL) ? sgl_memhs[elem_index]->uct[md_index] : - UCT_MEM_HANDLE_NULL; + iov.memh = (sgl_memhs != NULL) ? + ucp_datatype_iter_uct_memh(sgl_memhs[elem_index], + md_index) : + UCT_MEM_HANDLE_NULL; iov.stride = 0; iov.count = 1; diff --git a/test/gtest/ucp/test_ucp_dt.cc b/test/gtest/ucp/test_ucp_dt.cc index 8c55e59ad31..93e34664c54 100644 --- a/test/gtest/ucp/test_ucp_dt.cc +++ b/test/gtest/ucp/test_ucp_dt.cc @@ -430,13 +430,13 @@ class test_ucp_dt_sgl : public ucs::test { } private: - ucs::handle m_ucph; std::vector m_rkeys; ucp_rkey_t m_dummy_rkey; ucp_dt_local_sgl_t m_local; ucp_dt_remote_sgl_t m_remote; protected: + ucs::handle m_ucph; std::vector m_buffers; std::vector m_lengths; std::vector m_remote_addrs; @@ -452,6 +452,54 @@ class test_ucp_dt_sgl : public ucs::test { std::vector m_desc_elem_indices; }; +UCS_TEST_F(test_ucp_dt_sgl, mem_reg_provided_memh_missing_md) { + static constexpr size_t NUM_ELEMS = 2; + static constexpr size_t LENGTH = UCS_KBYTE; + static constexpr unsigned FLAGS = UCT_MD_MEM_ACCESS_RMA; + + ucp_md_map_t cache_md_map = m_ucph->cache_md_map[UCS_MEMORY_TYPE_HOST]; + if (ucs_popcount(cache_md_map) < 2) { + UCS_TEST_SKIP_R("need at least 2 cacheable memory domains"); + } + + /* Restricting md_map requires ucp_memh_get(), so these aren't user memhs */ + ucp_md_index_t md_index = ucs_ffs64(cache_md_map); + ucp_md_map_t reg_md_map = cache_md_map & ~UCS_BIT(md_index); + std::vector> data(NUM_ELEMS, + std::vector(LENGTH)); + std::vector buffers(NUM_ELEMS); + std::vector lengths(NUM_ELEMS, LENGTH); + std::vector memhs(NUM_ELEMS); + + for (size_t i = 0; i < NUM_ELEMS; ++i) { + buffers[i] = data[i].data(); + ASSERT_UCS_OK(ucp_memh_get(m_ucph, buffers[i], LENGTH, + UCS_MEMORY_TYPE_HOST, reg_md_map, FLAGS, + "test_sgl_memh", &memhs[i])); + ASSERT_FALSE(memhs[i]->md_map & UCS_BIT(md_index)); + } + + ucp_datatype_iter_t dt_iter = {}; + dt_iter.dt_class = UCP_DATATYPE_SGL; + dt_iter.length = NUM_ELEMS; + dt_iter.offset = 0; + dt_iter.type.sgl.buffers = buffers.data(); + dt_iter.type.sgl.lengths = lengths.data(); + dt_iter.type.sgl.memhs = memhs.data(); + dt_iter.type.sgl.frag_offset = 0; + ucp_memory_info_set_host(&dt_iter.mem_info); + + EXPECT_UCS_OK(ucp_datatype_iter_mem_reg(m_ucph, &dt_iter, + UCS_BIT(md_index), FLAGS, + UCS_BIT(UCP_DATATYPE_SGL))); + + for (size_t i = 0; i < NUM_ELEMS; ++i) { + EXPECT_TRUE(memhs[i]->md_map & UCS_BIT(md_index)); + EXPECT_TRUE(memhs[i]->uct[md_index] != UCT_MEM_HANDLE_NULL); + ucp_memh_put(memhs[i]); + } +} + UCS_TEST_F(test_ucp_dt_sgl, iter_next_chunked) { static constexpr size_t NUM_ELEMS = 10; static constexpr size_t MAX_PER_STEP = 3;