UCP: Do not use cuda staging if can't alloc cuda mem

openucx · Sep 20, 2024 · 2f4b402 · 2f4b402
1 parent 8ce6409
commit 2f4b402
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 13 deletions.
diff --git a/src/ucp/core/ucp_mm.c b/src/ucp/core/ucp_mm.c
@@ -1700,27 +1700,28 @@ ucs_status_t
 ucp_mm_get_alloc_md_index(ucp_context_h context, ucp_md_index_t *md_idx,
  ucs_memory_type_t alloc_mem_type)
 {
- ucs_status_t status;
+ ucs_status_t status = UCS_OK;
  uct_allocated_memory_t mem;
 
  if (!context->alloc_md[alloc_mem_type].initialized) {
+ context->alloc_md[alloc_mem_type].initialized = 1;
+
  status = ucp_mem_do_alloc(context, NULL, 1,
  UCT_MD_MEM_ACCESS_RMA |
  UCT_MD_MEM_FLAG_HIDE_ERRORS,
  alloc_mem_type, "get_alloc_md_id",
  &mem);
- if (status != UCS_OK) {
- return status;
+ if (status == UCS_OK) {
+ context->alloc_md[alloc_mem_type].md_index =
+ ucp_mem_get_md_index(context, mem.md, mem.method);
+ uct_mem_free(&mem);
+ } else {
+ context->alloc_md[alloc_mem_type].md_index = UCP_NULL_RESOURCE;
  }
-
- context->alloc_md[alloc_mem_type].initialized = 1;
- context->alloc_md[alloc_mem_type].md_index =
- ucp_mem_get_md_index(context, mem.md, mem.method);
- uct_mem_free(&mem);
  }
 
  *md_idx = context->alloc_md[alloc_mem_type].md_index;
- return UCS_OK;
+ return status;
 }
 
 static ucs_status_t

diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c
@@ -419,20 +419,24 @@ ucp_proto_rndv_rtr_mtype_probe(const ucp_proto_init_params_t *init_params)
  }
 
  params.mem_info.type = frag_mem_type;
+ params.md_map = 0;
 
  status = ucp_proto_init_buffer_copy_time(
  init_params->worker, "rtr/mtype unpack", frag_mem_type,
  init_params->select_param->mem_type, UCT_EP_OP_PUT_ZCOPY,
  &params.unpack_time, &params.unpack_perf_node);
  if (status != UCS_OK) {
- return;
+ continue;
  }
 
  status = ucp_mm_get_alloc_md_index(context, &md_index, frag_mem_type);
- if ((status != UCS_OK) || (md_index == UCP_NULL_RESOURCE)) {
- params.md_map = 0;
- } else {
+ if ((status == UCS_OK) && (md_index != UCP_NULL_RESOURCE)) {
  params.md_map = UCS_BIT(md_index);
+ } else if (frag_mem_type != UCS_MEMORY_TYPE_HOST) {
+ /* To use non-host staging buffers it should be possible to
+ * allocate them with MD */
+ ucp_proto_perf_node_deref(&params.unpack_perf_node);
+ continue;
  }
 
  rpriv.super.pack_cb = ucp_proto_rndv_rtr_mtype_pack;